diff --git a/README.md b/README.md index 4b2192b..87054df 100644 --- a/README.md +++ b/README.md @@ -74,6 +74,7 @@ Komandan has `komando` function that takes two arguments: - `ignore_exit_code`: a boolean that indicates whether to ignore the exit code of the task. If `true`, the script will continue even if the task returns a non-zero exit code. (default is `false`) - `elevate`: a boolean that indicates whether to run the task as root. (default is `false`) - `as_user`: a string that specifies the user to run the task as. Requires `elevate` to be `true`. (optional) + - `env`: a table that contains environment variables to be set before executing the task. (optional) This function will execute the module on the target server and return the results: - `stdout`: a string that contains the standard output of the module. @@ -83,11 +84,19 @@ This function will execute the module on the target server and return the result ## Modules Komandan has several built-in modules that can be used to perform various tasks on the target server. These modules are located in the `komandan.modules` table. + ### `cmd` module The `cmd` module allows you to execute a shell command on the target server. It takes the following arguments: - `cmd`: a string that contains the shell command to be executed. +Usage example: +```lua +local task = komandan.modules.cmd({ + cmd = "mkdir /tmp/newdir" +}) +``` + ### `script` module The `script` module allows you to execute a script on the target server. It takes the following arguments: @@ -95,18 +104,48 @@ The `script` module allows you to execute a script on the target server. It take - `from_file`: a string that contains the local path to the script file to be executed on the target server. (`script` and `from_file` parameters are mutually exclusive) - `interpreter`: a string that specifies the interpreter to use for the script. If not specified, the script will be executed using the default shell. +Usage example: + +```lua +local task = komandan.modules.script({ + script = "print('Hello from Komandan!')" + -- or + from_file = "/local_path/to/script.py" + + interpreter = "python3" +}) +``` + ### `upload` module The `upload` module allows you to upload a file to the target server. It takes the following arguments: - `src`: a string that contains the path to the file to be uploaded. - `dst`: a string that contains the path to the destination file on the target server. +Usage example: + +```lua +local task = komandan.modules.upload({ + src = "/local_path/to/file.txt", + dst = "/remote_path/to/file.txt" +}) +``` + ### `download` module The `download` module allows you to download a file from the target server. It takes the following arguments: - `src`: a string that contains the path to the file to be downloaded. - `dst`: a string that contains the path to the destination file on the local machine. +Usage example: + +```lua +local task = komandan.modules.download({ + src = "/remote_path/to/file.txt", + dst = "/local_path/to/file.txt" +}) +``` + ### `apt` module The `apt` module allows you to install packages on the target server. It takes the following arguments: @@ -115,6 +154,16 @@ The `apt` module allows you to install packages on the target server. It takes t - `update_cache`: a boolean that indicates whether to update the package cache before installing the package. (default is `false`) - `install_recommends`: a boolean that indicates whether to install recommended packages. (default is `true`) +Usage example: + +```lua +local task = komandan.modules.apt({ + package = "nginx", + update_cache = true, + elevate = true +}) +``` + ## Built-in functions Komandan provides several built-in functions that can be used to help write scripts. diff --git a/src/defaults.rs b/src/defaults.rs index 05c3081..033e017 100644 --- a/src/defaults.rs +++ b/src/defaults.rs @@ -8,5 +8,9 @@ pub fn defaults(lua: &Lua) -> mlua::Result { defaults.set("elevate", false)?; defaults.set("elevation_method", "sudo")?; + let env = lua.create_table()?; + env.set("DEBIAN_FRONTEND", "noninteractive")?; + defaults.set("env", env)?; + Ok(defaults) } diff --git a/src/main.rs b/src/main.rs index b1d689f..eaff330 100644 --- a/src/main.rs +++ b/src/main.rs @@ -252,13 +252,35 @@ async fn komando(lua: Lua, (host, task): (Value, Value)) -> mlua::Result
as_user, }; - let ssh = SSHSession::connect( + let mut ssh = SSHSession::connect( (host.get::("address")?.as_str(), port), &user, ssh_auth_method, elevation, )?; + let env_defaults = defaults.get::
("env").unwrap_or(lua.create_table()?); + let env_host = host.get::
("env").unwrap_or(lua.create_table()?); + let env_task = task.get::
("env").unwrap_or(lua.create_table()?); + + for pair in env_defaults.pairs() { + let (key, value): (String, String) = pair?; + println!("{}={}", key, value); + ssh.set_env(&key, &value); + } + + for pair in env_host.pairs() { + let (key, value): (String, String) = pair?; + println!("{}={}", key, value); + ssh.set_env(&key, &value); + } + + for pair in env_task.pairs() { + let (key, value): (String, String) = pair?; + println!("{}={}", key, value); + ssh.set_env(&key, &value); + } + let module_clone = module.clone(); let results = lua .load(chunk! { @@ -285,11 +307,6 @@ async fn komando(lua: Lua, (host, task): (Value, Value)) -> mlua::Result
}) .eval::<()>()?; - let defaults = lua - .globals() - .get::
("komandan")? - .get::
("defaults")?; - let ignore_exit_code = task .get::("ignore_exit_code") .unwrap_or_else(|_| defaults.get::("ignore_exit_code").unwrap()); diff --git a/src/ssh.rs b/src/ssh.rs index 3584407..7e9fd9a 100644 --- a/src/ssh.rs +++ b/src/ssh.rs @@ -1,4 +1,5 @@ use std::{ + collections::HashMap, fs, io::{self, Read, Write}, net::{TcpStream, ToSocketAddrs}, @@ -30,6 +31,7 @@ pub enum ElevateMethod { pub struct SSHSession { session: Session, + env: HashMap, elevation: Elevation, stdout: Option, stderr: Option, @@ -72,6 +74,7 @@ impl SSHSession { Ok(Self { session, + env: HashMap::new(), elevation, stdout: Some(String::new()), stderr: Some(String::new()), @@ -79,17 +82,30 @@ impl SSHSession { }) } + pub fn set_env(&mut self, key: &str, value: &str) { + *self.env.entry(key.to_string()).or_insert(value.to_string()) = value.to_string(); + } + + fn execute_command(&mut self, command: &str) -> Result { + let mut channel = self.session.channel_session()?; + let mut command = command.to_string(); + for (key, value) in &self.env { + command = format!("export {}={}\n", key, value) + &command; + } + channel.exec(command.as_str())?; + Ok(channel) + } + pub fn cmd(&mut self, command: &str) -> Result<(String, String, i32)> { - let mut channel = self.session.channel_session().unwrap(); - channel.exec(command).unwrap(); + let mut channel = self.execute_command(command)?; let mut stdout = String::new(); let mut stderr = String::new(); - channel.read_to_string(&mut stdout).unwrap(); - channel.stderr().read_to_string(&mut stderr).unwrap(); + channel.read_to_string(&mut stdout)?; + channel.stderr().read_to_string(&mut stderr)?; stdout = stdout.trim_end_matches('\n').to_string(); - channel.wait_close().unwrap(); - let exit_code = channel.exit_status().unwrap(); + channel.wait_close()?; + let exit_code = channel.exit_status()?; self.stdout.as_mut().unwrap().push_str(&stdout); self.stderr.as_mut().unwrap().push_str(&stderr); @@ -115,32 +131,27 @@ impl SSHSession { } pub fn get_remote_env(&mut self, var: &str) -> Result { - let mut channel = self.session.channel_session().unwrap(); - channel.exec(format!("echo ${}", var).as_str()).unwrap(); + let mut channel = self.execute_command(format!("echo ${}", var).as_str())?; let mut stdout = String::new(); - channel.read_to_string(&mut stdout).unwrap(); + channel.read_to_string(&mut stdout)?; stdout = stdout.trim_end_matches('\n').to_string(); - channel.wait_close().unwrap(); + channel.wait_close()?; Ok(stdout) } pub fn get_tmpdir(&mut self) -> Result { - let mut channel = self.session.channel_session().unwrap(); - channel.exec("tmpdir=`for dir in \"$HOME/.komandan/tmp\" \"/tmp/komandan\"; do if [ -d \"$dir\" ] || mkdir -p \"$dir\" 2>/dev/null; then echo \"$dir\"; break; fi; done`; [ -z \"$tmpdir\" ] && { exit 1; } || echo \"$tmpdir\"").unwrap(); + let mut channel = self.execute_command("tmpdir=`for dir in \"$HOME/.komandan/tmp\" \"/tmp/komandan\"; do if [ -d \"$dir\" ] || mkdir -p \"$dir\" 2>/dev/null; then echo \"$dir\"; break; fi; done`; [ -z \"$tmpdir\" ] && { exit 1; } || echo \"$tmpdir\"")?; let mut stdout = String::new(); - channel.read_to_string(&mut stdout).unwrap(); + channel.read_to_string(&mut stdout)?; stdout = stdout.trim_end_matches('\n').to_string(); - channel.wait_close().unwrap(); + channel.wait_close()?; Ok(stdout) } pub fn chmod(&mut self, remote_path: &Path, mode: &String) -> Result<()> { - let mut channel = self.session.channel_session().unwrap(); - channel - .exec(format!("chmod {} {}", mode, remote_path.to_string_lossy()).as_str()) - .unwrap(); + self.execute_command(format!("chmod {} {}", mode, remote_path.to_string_lossy()).as_str())?; Ok(()) } @@ -171,15 +182,14 @@ impl SSHSession { pub fn write_remote_file(&mut self, remote_path: &str, content: &[u8]) -> Result<()> { let content_length = content.len() as u64; - let mut remote_file = self - .session - .scp_send(Path::new(remote_path), 0o644, content_length, None) - .unwrap(); - remote_file.write(content).unwrap(); - remote_file.send_eof().unwrap(); - remote_file.wait_eof().unwrap(); - remote_file.close().unwrap(); - remote_file.wait_close().unwrap(); + let mut remote_file = + self.session + .scp_send(Path::new(remote_path), 0o644, content_length, None)?; + remote_file.write(content)?; + remote_file.send_eof()?; + remote_file.wait_eof()?; + remote_file.close()?; + remote_file.wait_close()?; Ok(()) } @@ -251,9 +261,9 @@ fn download_directory(sftp: &mut Sftp, remote_path: &Path, local_path: &Path) -> impl UserData for SSHSession { fn add_methods>(methods: &mut M) { methods.add_method_mut("cmd", |lua, this, command: String| { - let command = this.prepare_command(command.as_str()).unwrap(); + let command = this.prepare_command(command.as_str())?; let cmd_result = this.cmd(&command); - let (stdout, stderr, exit_code) = cmd_result.unwrap(); + let (stdout, stderr, exit_code) = cmd_result?; let table = lua.create_table()?; table.set("stdout", stdout)?; @@ -299,7 +309,7 @@ impl UserData for SSHSession { }); methods.add_method_mut("get_tmpdir", |_, this, ()| { - let tmpdir = this.get_tmpdir().unwrap(); + let tmpdir = this.get_tmpdir()?; Ok(tmpdir) });