From 7b9066d096cfbb98ec825dc647e9e10e9a235627 Mon Sep 17 00:00:00 2001 From: "R. Tyler Croy" Date: Wed, 30 Dec 2020 21:11:35 -0800 Subject: [PATCH] Implement the simple support for running a task This includes an echo task for funsies --- Cargo.lock | 117 +++++++++++++++++++++++++++++++ cli/Cargo.toml | 3 + cli/src/main.rs | 108 +++++++++++++++++++++++----- cli/src/transport/mod.rs | 8 ++- cli/src/transport/ssh.rs | 19 +++-- native-tasks/.gitignore | 0 parser/Cargo.toml | 1 + parser/src/lib.rs | 147 ++++++++++++++++++++++++++++++++++++++- parser/src/task.pest | 10 +-- tasks/.gitignore | 0 tasks/echo.ztask | 12 ++++ 11 files changed, 397 insertions(+), 28 deletions(-) create mode 100644 native-tasks/.gitignore create mode 100644 tasks/.gitignore create mode 100644 tasks/echo.ztask diff --git a/Cargo.lock b/Cargo.lock index 12876eb..ac2eb74 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1,5 +1,14 @@ # This file is automatically @generated by Cargo. # It is not intended for manual editing. +[[package]] +name = "aho-corasick" +version = "0.7.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7404febffaa47dac81aa44dba71523c9d069b1bdc50a77db41195149e17f68e5" +dependencies = [ + "memchr", +] + [[package]] name = "atty" version = "0.2.14" @@ -103,6 +112,19 @@ version = "0.4.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "88d7ed2934d741c6b37e33e3832298e8850b53fd2d2bea03873375596c7cea4e" +[[package]] +name = "env_logger" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44533bbbb3bb3c1fa17d9f2e4e38bbbaf8396ba82193c4cb1b6445d711445d36" +dependencies = [ + "atty", + "humantime", + "log", + "regex", + "termcolor", +] + [[package]] name = "fake-simd" version = "0.1.2" @@ -118,6 +140,12 @@ dependencies = [ "typenum", ] +[[package]] +name = "glob" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b919933a397b79c37e33b77bb2aa3dc8eb6e165ad809e58ff75bc7db2e34574" + [[package]] name = "gumdrop" version = "0.8.0" @@ -147,6 +175,15 @@ dependencies = [ "libc", ] +[[package]] +name = "humantime" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df004cfca50ef23c36850aaaa59ad52cc70d0e90243c3c7737a4dd32dc7a3c4f" +dependencies = [ + "quick-error", +] + [[package]] name = "itoa" version = "0.4.7" @@ -206,12 +243,27 @@ dependencies = [ "scopeguard", ] +[[package]] +name = "log" +version = "0.4.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fabed175da42fed1fa0746b0ea71f412aa9d35e76e95e59b192c64b9dc2bf8b" +dependencies = [ + "cfg-if", +] + [[package]] name = "maplit" version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3e2e65a1a2e43cfcb47a895c4c8b10d1f4a61097f9f254f183aee60cad9c651d" +[[package]] +name = "memchr" +version = "2.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ee1c47aaa256ecabcaea351eae4a9b01ef39ed810004e298d2511ed284b1525" + [[package]] name = "opaque-debug" version = "0.2.3" @@ -304,6 +356,16 @@ version = "0.3.19" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3831453b3449ceb48b6d9c7ad7c96d5ea673e9b470a1dc578c2ce6521230884c" +[[package]] +name = "pretty_env_logger" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "926d36b9553851b8b0005f1275891b392ee4d2d833852c417ed025477350fb9d" +dependencies = [ + "env_logger", + "log", +] + [[package]] name = "proc-macro2" version = "1.0.24" @@ -313,6 +375,12 @@ dependencies = [ "unicode-xid", ] +[[package]] +name = "quick-error" +version = "1.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a1d01941d82fa2ab50be1e79e6714289dd7cde78eba4c074bc5a4374f650dfe0" + [[package]] name = "quote" version = "1.0.8" @@ -328,6 +396,24 @@ version = "0.1.57" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "41cc0f7e4d5d4544e8861606a285bb08d3e70712ccc7d2b84d7c0ccfaf4b05ce" +[[package]] +name = "regex" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38cf2c13ed4745de91a5eb834e11c00bcc3709e773173b2ce4c56c9fbde04b9c" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", + "thread_local", +] + +[[package]] +name = "regex-syntax" +version = "0.6.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b181ba2dcf07aaccad5448e8ead58db5b742cf85dfe035e2227f137a539a189" + [[package]] name = "ryu" version = "1.0.5" @@ -424,6 +510,24 @@ dependencies = [ "unicode-xid", ] +[[package]] +name = "termcolor" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2dfed899f0eb03f32ee8c6a0aabdb8a7949659e3466561fc0adf54e26d88c5f4" +dependencies = [ + "winapi-util", +] + +[[package]] +name = "thread_local" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d40c6d1b69745a6ec6fb1ca717914848da4b44ae29d9b3080cbee91d72a69b14" +dependencies = [ + "lazy_static", +] + [[package]] name = "typenum" version = "1.12.0" @@ -464,6 +568,15 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" +[[package]] +name = "winapi-util" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "70ec6ce85bb158151cae5e5c87f95a8e97d2c0c4b001223f33a334e3ce5de178" +dependencies = [ + "winapi", +] + [[package]] name = "winapi-x86_64-pc-windows-gnu" version = "0.4.0" @@ -484,7 +597,10 @@ name = "zap" version = "0.1.0" dependencies = [ "colored", + "glob", "gumdrop", + "log", + "pretty_env_logger", "serde", "serde_derive", "serde_json", @@ -497,6 +613,7 @@ dependencies = [ name = "zap-parser" version = "0.1.0" dependencies = [ + "log", "pest", "pest_derive", ] diff --git a/cli/Cargo.toml b/cli/Cargo.toml index 30c64a4..77bb6f2 100644 --- a/cli/Cargo.toml +++ b/cli/Cargo.toml @@ -6,7 +6,10 @@ edition = "2018" [dependencies] colored = "2" +glob = "0.3" gumdrop = "~0.8.0" +log = "0.4" +pretty_env_logger = "0.4" # Needed for deserializing JSON messages _and_ managing our configuration # effectively serde = { version = "~1.0", features = ["derive", "rc"] } diff --git a/cli/src/main.rs b/cli/src/main.rs index 0286013..9b86c46 100644 --- a/cli/src/main.rs +++ b/cli/src/main.rs @@ -1,15 +1,18 @@ use colored::*; use gumdrop::Options; +use log::*; use std::io::BufReader; mod inventory; mod transport; +use zap_parser::*; use crate::inventory::*; use crate::transport::ssh::Ssh; use crate::transport::Transport; fn main() { + pretty_env_logger::init(); let opts = MyOptions::parse_args_default_or_exit(); if opts.command.is_none() { @@ -27,22 +30,86 @@ fn main() { }; match opts.command.unwrap() { - Command::Cmd(runopts) => { - if let Some(group) = inventory.groups.iter().find(|g| g.name == runopts.targets) { - std::process::exit(runner.run_group(&runopts.command, &group, &inventory)); - } - - if let Some(target) = inventory.targets.iter().find(|t| t.name == runopts.targets) { - println!("{}", format!("run a command: {:?}", runopts).green()); - std::process::exit(runner.run(&runopts.command, &target)); - } - - println!("{}", format!("Couldn't find a target named `{}`", runopts.targets).red()); - } + Command::Cmd(opts) => handle_cmd(opts, &runner, inventory), + Command::Task(opts) => handle_task(opts, &runner, inventory), _ => {} } } +fn load_ztasks() -> Vec { + use glob::glob; + let mut tasks = vec![]; + + for entry in glob("tasks/**/*.ztask").expect("Failed to read glob pattern") { + match entry { + Ok(path) => { + let task = Task::from_path(&path).unwrap(); + if let Ok(task) = Task::from_path(&path) { + info!("loaded ztask: {}", task.name); + tasks.push(task); + } + }, + Err(e) => println!("{:?}", e), + } + } + tasks +} + +/** + * This function will handle a task + */ +fn handle_task(opts: TaskOpts, runner: &dyn crate::transport::Transport, inventory: Inventory) { + println!("running task: {:?}", opts); + + for task in load_ztasks() { + if task.name == opts.task { + let mut env = crate::transport::EnvVars::new(); + + for parameter in opts.parameter.iter() { + let parts: Vec<&str> = parameter.split("=").collect(); + if parts.len() == 2 { + env.insert(parts[0].to_string(), parts[1].to_string()); + } + } + + if let Some(script) = task.get_script() { + // TODO: refactor with handle_cmd + if let Some(group) = inventory.groups.iter().find(|g| g.name == opts.targets) { + std::process::exit(runner.run_group(&script, &group, &inventory, Some(env))); + } + + if let Some(target) = inventory.targets.iter().find(|t| t.name == opts.targets) { + std::process::exit(runner.run(&script, &target, Some(&env))); + } + + } + } + } +} + +/** + * This function will handle executing a single specified command on the target(s) + * identified in the `opts`. + * + * In the case of a single target, the status code from the executed command will + * be propogated up. + * + * In the case of multiple targets, any non-zero status code will be used to exit + * non-zero. + */ +fn handle_cmd(opts: CmdOpts, runner: &dyn crate::transport::Transport, inventory: Inventory) { + if let Some(group) = inventory.groups.iter().find(|g| g.name == opts.targets) { + std::process::exit(runner.run_group(&opts.command, &group, &inventory, None)); + } + + if let Some(target) = inventory.targets.iter().find(|t| t.name == opts.targets) { + println!("{}", format!("run a command: {:?}", opts).green()); + std::process::exit(runner.run(&opts.command, &target, None)); + } + + println!("{}", format!("Couldn't find a target named `{}`", opts.targets).red()); +} + #[derive(Debug, Options)] struct MyOptions { // Options here can be accepted with any command (or none at all), @@ -73,7 +140,9 @@ enum Command { #[options(help = "show help for a command")] Help(HelpOpts), #[options(help = "Run a single command on a target(s)")] - Cmd(RunOpts), + Cmd(CmdOpts), + #[options(help = "Execute a task on a target(s)")] + Task(TaskOpts), } #[derive(Debug, Options)] @@ -81,12 +150,19 @@ struct HelpOpts { #[options(free)] free: Vec, } - -// Options accepted for the `make` command #[derive(Debug, Options)] -struct RunOpts { +struct CmdOpts { #[options(free, help = "Command to execute on the target(s)")] command: String, #[options(help = "Name of a target or group")] targets: String, } +#[derive(Debug, Options)] +struct TaskOpts { + #[options(free, help = "Task to execute, must exist in ZAP_PATH")] + task: String, + #[options(short="p", help = "Parameter values")] + parameter: Vec, + #[options(help = "Name of a target or group")] + targets: String, +} diff --git a/cli/src/transport/mod.rs b/cli/src/transport/mod.rs index 4ef4c18..f4fe68d 100644 --- a/cli/src/transport/mod.rs +++ b/cli/src/transport/mod.rs @@ -1,4 +1,8 @@ use crate::inventory::{Group, Inventory, Target}; +use std::collections::HashMap; + + +pub type EnvVars = HashMap; pub mod ssh; @@ -7,6 +11,6 @@ pub mod ssh; * connecting to targets */ pub trait Transport { - fn run_group(&self, cmd: &str, group: &Group, inv: &Inventory) -> i32; - fn run(&self, command: &str, target: &Target) -> i32; + fn run_group(&self, cmd: &str, group: &Group, inv: &Inventory, env: Option) -> i32; + fn run(&self, command: &str, target: &Target, env: Option<&EnvVars>) -> i32; } diff --git a/cli/src/transport/ssh.rs b/cli/src/transport/ssh.rs index 99b771b..010be81 100644 --- a/cli/src/transport/ssh.rs +++ b/cli/src/transport/ssh.rs @@ -1,4 +1,5 @@ use crate::inventory::{Group, Inventory, Target}; +use crate::transport::EnvVars; use crate::transport::Transport; use serde::{Deserialize, Serialize}; @@ -16,20 +17,20 @@ impl Default for Ssh { } impl Transport for Ssh { - fn run_group(&self, command: &str, group: &Group, inventory: &Inventory) -> i32 { + fn run_group(&self, command: &str, group: &Group, inventory: &Inventory, env: Option) -> i32 { let mut status = 1; for target_name in group.targets.iter() { // XXX: This is inefficient for target in inventory.targets.iter() { if &target.name == target_name { println!("Running on `{}`", target.name); - status = self.run(command, &target); + status = self.run(command, &target, env.as_ref()); } } } status } - fn run(&self, command: &str, target: &Target) -> i32 { + fn run(&self, command: &str, target: &Target, env: Option<&EnvVars>) -> i32 { // Connect to the local SSH server let tcp = TcpStream::connect(format!("{}:22", target.uri)).unwrap(); let mut sess = Session::new().unwrap(); @@ -39,7 +40,17 @@ impl Transport for Ssh { .unwrap(); let mut channel = sess.channel_session().unwrap(); - channel.exec(command).unwrap(); + + let mut segments = vec![]; + + if let Some(env) = env { + for (key, val) in env.iter() { + segments.push(format!("export ZAP_{}=\"{}\"", key.to_uppercase(), val)); + } + } + segments.push(command.to_string()); + + channel.exec(&segments.join(";")).unwrap(); let mut s = String::new(); channel.read_to_string(&mut s).unwrap(); print!("{}", s); diff --git a/native-tasks/.gitignore b/native-tasks/.gitignore new file mode 100644 index 0000000..e69de29 diff --git a/parser/Cargo.toml b/parser/Cargo.toml index 429b9ad..d9dad81 100644 --- a/parser/Cargo.toml +++ b/parser/Cargo.toml @@ -5,5 +5,6 @@ authors = ["R. Tyler Croy "] edition = "2018" [dependencies] +log = "0.4" pest = "~2.1" pest_derive = "~2.1" diff --git a/parser/src/lib.rs b/parser/src/lib.rs index bf880fd..a931bdd 100644 --- a/parser/src/lib.rs +++ b/parser/src/lib.rs @@ -4,11 +4,144 @@ extern crate pest; extern crate pest_derive; use pest::Parser; +use pest::error::Error as PestError; +use pest::error::ErrorVariant; +use pest::iterators::Pairs; +use std::path::PathBuf; #[derive(Parser)] #[grammar="task.pest"] struct TaskParser; + +pub struct Task { + pub name: String, + inline: Option, +} + +impl Task { + pub fn get_script(&self) -> Option<&String> { + self.inline.as_ref() + } + + pub fn new(name: &str) -> Self { + Task { + name: name.to_string(), + inline: None, + } + } + + fn parse(parser: &mut Pairs) -> Result> { + let mut task: Option = None; + + while let Some(parsed) = parser.next() { + match parsed.as_rule() { + Rule::identifier => { + task = Some(Task::new(parsed.as_str())); + }, + Rule::script => { + let script = parse_str(&mut parsed.into_inner())?; + + if let Some(ref mut task) = task { + task.inline = Some(script); + } + } + _ => {}, + } + } + + if let Some(task) = task { + return Ok(task); + } + else { + return Err(PestError::new_from_pos( + ErrorVariant::CustomError { + message: "Could not find a valid task definition".to_string(), + }, + /* TODO: Find a better thing to report */ + pest::Position::from_start(""), + )); + } + } + + pub fn from_str(buf: &str) -> Result> { + let mut parser = TaskParser::parse(Rule::task, buf)?; + while let Some(parsed) = parser.next() { + match parsed.as_rule() { + Rule::task => { + return Task::parse(&mut parsed.into_inner()); + }, + _ => {}, + } + } + return Err(PestError::new_from_pos( + ErrorVariant::CustomError { + message: "Could not find a valid task definition".to_string(), + }, + pest::Position::from_start(buf), + )); + } + + pub fn from_path(path: &PathBuf) -> Result> { + use std::fs::File; + use std::io::Read; + + match File::open(path) { + Ok(mut file) => { + let mut contents = String::new(); + + if let Err(e) = file.read_to_string(&mut contents) { + return Err(PestError::new_from_pos( + ErrorVariant::CustomError { + message: format!("{}", e), + }, + pest::Position::from_start(""), + )); + } else { + return Self::from_str(&contents); + } + } + Err(e) => { + return Err(PestError::new_from_pos( + ErrorVariant::CustomError { + message: format!("{}", e), + }, + pest::Position::from_start(""), + )); + } + } + } +} + +/** + * Parser utility function to fish out the _actual_ string value for something + * that is looking like a string Rule + * + */ +fn parse_str(parser: &mut Pairs) -> Result> { + while let Some(parsed) = parser.next() { + match parsed.as_rule() { + Rule::string => { + return parse_str(&mut parsed.into_inner()); + }, + Rule::double_quoted => { + return parse_str(&mut parsed.into_inner()); + }, + Rule::inner_double_str => { + return Ok(parsed.as_str().to_string()); + } + _ => {} + } + } + return Err(PestError::new_from_pos( + ErrorVariant::CustomError { + message: "Could not parse out a string value".to_string(), + }, + /* TODO: Find a better thing to report */ + pest::Position::from_start(""), + )); +} + #[cfg(test)] mod tests { use super::*; @@ -38,7 +171,19 @@ mod tests { inline = "env" } }"#; - let _task = TaskParser::parse(Rule::task, buf) + let task = TaskParser::parse(Rule::task, buf) .unwrap().next().unwrap(); } + + #[test] + fn parse_task_fn() { + let buf = r#"task PrintEnv { + script { + inline = "env" + } + }"#; + let task = Task::from_str(buf).expect("Failed to parse the task"); + assert_eq!(task.name, "PrintEnv"); + assert_eq!(task.get_script().unwrap(), "env"); + } } diff --git a/parser/src/task.pest b/parser/src/task.pest index 641371d..830eef5 100644 --- a/parser/src/task.pest +++ b/parser/src/task.pest @@ -38,12 +38,12 @@ script = { "script" ~ (script_inline) ~ closing_brace } -script_inline = { "inline" ~ equals ~ string } +script_inline = _{ "inline" ~ equals ~ string } -opening_brace = { "{" } -closing_brace = { "}" } -equals = { "=" } -quote = { "\"" } +opening_brace = _{ "{" } +closing_brace = _{ "}" } +equals = _{ "=" } +quote = _{ "\"" } string = { double_quoted } double_quoted = ${ (quote ~ inner_double_str ~ quote) } diff --git a/tasks/.gitignore b/tasks/.gitignore new file mode 100644 index 0000000..e69de29 diff --git a/tasks/echo.ztask b/tasks/echo.ztask new file mode 100644 index 0000000..41482a0 --- /dev/null +++ b/tasks/echo.ztask @@ -0,0 +1,12 @@ +task Echo { + parameters { + msg { + required = true + help = "String to echo back to the client" + type = string + } + } + script { + inline = "echo ${ZAP_MSG}" + } +}