rtyler
/
zap
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Refactor the SSH transport to make it easy to re-send files 

This also has some refactor to ensure that a set up transport stays connected
instead of reconnecting for each "run"
This commit is contained in:
R Tyler Croy 2020-12-31 21:53:19 -08:00
parent 7e3d7fb5ad
commit bc4396a0de
3 changed files with 92 additions and 64 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
cli/src/main.rs
View File

@ -27,14 +27,14 @@ fn main() {
let reader = BufReader::new(file);
let inventory: Inventory = serde_yaml::from_reader(reader).expect("Failed to read intenvory");
let runner = match &inventory.config.transport {
let mut runner = match &inventory.config.transport {
crate::inventory::Transport::Ssh => Ssh::default(),
};
match opts.command.unwrap() {
Command::Cmd(opts) => handle_cmd(opts, &runner, inventory),
Command::Task(opts) => handle_task(opts, &runner, inventory),
Command::Plan(opts) => handle_plan(opts, &runner, inventory),
Command::Cmd(opts) => handle_cmd(opts, &mut runner, inventory),
Command::Task(opts) => handle_task(opts, &mut runner, inventory),
Command::Plan(opts) => handle_plan(opts, &mut runner, inventory),
Command::Check(opts) => handle_check(opts),
_ => {}
}
@ -62,7 +62,7 @@ fn handle_check(opts: CheckOpts) {
/**
* This function will parse and execute a plan
*/
fn handle_plan(opts: PlanOpts, runner: &dyn crate::transport::Transport, inventory: Inventory) {
fn handle_plan(opts: PlanOpts, runner: &mut dyn crate::transport::Transport, inventory: Inventory) {
println!("{}", format!("Running plan with: {:?}", opts).green());
let mut exit: i32 = -1;
@ -90,7 +90,7 @@ fn handle_plan(opts: PlanOpts, runner: &dyn crate::transport::Transport, invento
fn execute_task_on(
targets: String,
task: &ExecutableTask,
runner: &dyn crate::transport::Transport,
runner: &mut dyn crate::transport::Transport,
inventory: &Inventory,
dry_run: bool,
) -> i32 {
@ -108,7 +108,7 @@ fn execute_task_on(
/**
* This function will handle a task
*/
fn handle_task(opts: TaskOpts, runner: &dyn crate::transport::Transport, inventory: Inventory) {
fn handle_task(opts: TaskOpts, runner: &mut dyn crate::transport::Transport, inventory: Inventory) {
println!("{}", format!("Running task with: {:?}", opts).green());
match Task::from_path(&opts.task) {
@ -153,7 +153,7 @@ fn handle_task(opts: TaskOpts, runner: &dyn crate::transport::Transport, invento
* 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) {
fn handle_cmd(opts: CmdOpts, runner: &mut dyn crate::transport::Transport, inventory: Inventory) {
let mut task = ExecutableTask::new(Task::new("Dynamic"), HashMap::new());
task.task.script.inline = Some(opts.command);
std::process::exit(execute_task_on(

14
cli/src/transport/mod.rs
View File

@ -1,4 +1,5 @@
use crate::inventory::{Group, Inventory, Target};
use std::path::Path;
use zap_parser::plan::ExecutableTask;
pub mod ssh;
@ -8,7 +9,14 @@ pub mod ssh;
* connecting to targets
*/
pub trait Transport {
fn run_group(&self, cmd: &ExecutableTask, group: &Group, inv: &Inventory, dry_run: bool)
-> i32;
fn run(&self, command: &ExecutableTask, target: &Target, dry_run: bool) -> i32;
fn connect(&mut self, target: &Target) -> bool;
fn run_group(
&mut self,
cmd: &ExecutableTask,
group: &Group,
inv: &Inventory,
dry_run: bool,
) -> i32;
fn run(&mut self, command: &ExecutableTask, target: &Target, dry_run: bool) -> i32;
fn send_bytes(&self, remote_path: &Path, bytes: &Vec<u8>, mode: i32) -> bool;
}

126
cli/src/transport/ssh.rs
View File

@ -3,7 +3,6 @@ use crate::transport::Transport;
use colored::*;
use log::*;
use serde::{Deserialize, Serialize};
use ssh2::Session;
use std::convert::TryInto;
use std::io::prelude::*;
@ -13,18 +12,24 @@ use std::path::Path;
use zap_parser::plan::ExecutableTask;
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct Ssh {}
#[derive(Clone)]
pub struct Ssh {
session: Session,
connected: bool,
}
impl Default for Ssh {
fn default() -> Self {
Self {}
Self {
session: Session::new().unwrap(),
connected: false,
}
}
}
impl Transport for Ssh {
fn run_group(
&self,
&mut self,
command: &ExecutableTask,
group: &Group,
inventory: &Inventory,
@ -43,30 +48,40 @@ impl Transport for Ssh {
status
}
fn run(&self, command: &ExecutableTask, target: &Target, dry_run: bool) -> i32 {
// Connect to the local SSH server
fn connect(&mut self, target: &Target) -> bool {
if self.connected {
return self.connected;
}
let tcp = TcpStream::connect(format!("{}:22", target.uri)).unwrap();
let mut sess = Session::new().unwrap();
sess.set_tcp_stream(tcp);
sess.handshake().unwrap();
self.session.set_tcp_stream(tcp);
self.session.handshake().unwrap();
let mut authenticated = false;
if let Some(config) = &target.config {
if let Some(sshconfig) = &config.ssh {
// requires PasswordAuthentication yes
sess.userauth_password(&sshconfig.user, &sshconfig.password)
self.session
.userauth_password(&sshconfig.user, &sshconfig.password)
.unwrap();
authenticated = true;
}
}
if !authenticated {
sess.userauth_agent(&std::env::var("USER").unwrap())
self.session
.userauth_agent(&std::env::var("USER").unwrap())
.unwrap();
}
let remote_script = "._zap_command";
let args_file = "._zap_args.json";
self.connected = true;
true
}
fn run(&mut self, command: &ExecutableTask, target: &Target, dry_run: bool) -> i32 {
if !self.connect(target) {
error!("Failed to connect to {:?}", target);
return -1;
}
if let Some(provides) = &command.parameters.get("provides") {
debug!(
@ -74,7 +89,7 @@ impl Transport for Ssh {
provides
);
if let Err(error) = sess.scp_recv(&Path::new(&provides)) {
if let Err(error) = self.session.scp_recv(&Path::new(&provides)) {
if error.code() == ssh2::ErrorCode::Session(-28) {
debug!(
"The provided file ({}) does not exist, the command should be run",
@ -98,6 +113,13 @@ impl Transport for Ssh {
}
}
if let Some(unless) = &command.parameters.get("unless") {
debug!("An `unless` parameter was given, running {}", unless);
}
let remote_script = "._zap_command";
let args_file = "._zap_args.json";
if let Some(script) = command.task.script.as_bytes(Some(&command.parameters)) {
if dry_run {
println!("{}", "Dry-run\n----".yellow());
@ -108,49 +130,25 @@ impl Transport for Ssh {
return 0;
}
let mut remote_file = sess
.scp_send(
Path::new(remote_script),
0o700,
script
.len()
.try_into()
.expect("Overflow converting the size of the generated file, yikes!"),
None,
)
.unwrap();
remote_file.write(&script).unwrap();
// Close the channel and wait for the whole content to be tranferred
remote_file.send_eof().unwrap();
remote_file.wait_eof().unwrap();
remote_file.close().unwrap();
remote_file.wait_close().unwrap();
if !self.send_bytes(Path::new(remote_script), &script, 0o700) {
error!("Failed to upload script file for execution");
return -1;
}
let mut channel = sess.channel_session().unwrap();
let mut channel = self.session.channel_session().unwrap();
let stderr = channel.stderr();
if command.task.script.has_file() {
let args = serde_json::to_string(&command.parameters)
.expect("Failed to serialize parameters for task");
let mut remote_file = sess
.scp_send(
Path::new(args_file),
0o400,
args.len().try_into().expect(
"Failed converting the size of the generated args file, yikes!",
),
None,
)
.unwrap();
remote_file.write(&args.as_bytes()).unwrap();
// Close the channel and wait for the whole content to be tranferred
remote_file.send_eof().unwrap();
remote_file.wait_eof().unwrap();
remote_file.close().unwrap();
remote_file.wait_close().unwrap();
channel
.exec(&format!("./{} {}", remote_script, args_file))
.unwrap();
if self.send_bytes(Path::new(args_file), &args.into_bytes(), 0o400) {
channel
.exec(&format!("./{} {}", remote_script, args_file))
.unwrap();
} else {
error!("Failed to upload the arguments file");
return -1;
}
} else {
channel.exec(&format!("./{}", remote_script)).unwrap();
}
@ -171,7 +169,7 @@ impl Transport for Ssh {
* somehow and I'm not seeing anything that would allow me to just reach
* out and remove a file
*/
let mut channel = sess.channel_session().unwrap();
let mut channel = self.session.channel_session().unwrap();
channel
.exec(&format!("rm -f {} {}", remote_script, args_file))
.unwrap();
@ -181,4 +179,26 @@ impl Transport for Ssh {
return -1;
}
}
fn send_bytes(&self, remote_path: &Path, bytes: &Vec<u8>, mode: i32) -> bool {
let mut remote_file = self
.session
.scp_send(
remote_path,
mode,
bytes
.len()
.try_into()
.expect("Failed converting the size of the file to send, yikes!"),
None,
)
.unwrap();
remote_file.write(bytes).unwrap();
// Close the channel and wait for the whole content to be tranferred
remote_file.send_eof().unwrap();
remote_file.wait_eof().unwrap();
remote_file.close().unwrap();
remote_file.wait_close().unwrap();
true
}
}