add new modules and update readme

README.md

@@ -1,15 +1,16 @@
 <div align="center">
 
 # Komandan
-##### Your army commander
+#### Your army commander
 
-[![Build Status]][github-actions] [![License:MIT]][license]
+[![Build Status]][github-actions] [![License:MIT]][license] [![Coverage]][codecov.io]
 
 [Build Status]: https://github.com/hahnavi/komandan/actions/workflows/rust.yml/badge.svg
 [github-actions]: https://github.com/hahnavi/komandan/actions
-[License:MIT]: https://img.shields.io/badge/License-MIT-yellow.svg
-[license]: https://github.com/hahnavi/komandan/blob/master/LICENSE
-
+[License:MIT]: https://img.shields.io/badge/License-MIT-blue.svg
+[license]: https://github.com/hahnavi/komandan/blob/main/LICENSE
+[Coverage]: https://codecov.io/gh/hahnavi/komandan/branch/main/graph/badge.svg
+[codecov.io]: https://app.codecov.io/gh/hahnavi/komandan
 
 </div>
 
@@ -72,3 +73,15 @@ The `cmd` module allows you to execute a shell command on the target server. It
 The `script` module allows you to execute a script on the target server. It takes the following arguments:
 - `script`: a string that contains the script to be executed.
 - `interpreter`: a string that specifies the interpreter to use for the script. If not specified, the script will be executed using the default shell.
+
+### `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.
+
+### `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.

src/main.rs

@@ -1,7 +1,7 @@
 use args::Args;
 use clap::Parser;
 use mlua::{chunk, Error::RuntimeError, Integer, Lua, Table, Value};
-use modules::{cmd, script};
+use modules::{cmd, script, upload, download};
 use ssh::SSHSession;
 use std::path::Path;
 
@@ -110,6 +110,8 @@ fn setup_komandan_table(lua: &Lua) -> mlua::Result<()> {
     let modules_table = lua.create_table()?;
     modules_table.set("cmd", lua.create_function(cmd)?)?;
     modules_table.set("script", lua.create_function(script)?)?;
+    modules_table.set("upload", lua.create_function(upload)?)?;
+    modules_table.set("download", lua.create_function(download)?)?;
     komandan.set("modules", modules_table)?;
 
     lua.globals().set("komandan", &komandan)?;

src/modules/download.rs

@@ -0,0 +1,20 @@
+use mlua::{chunk, ExternalResult, Lua, Table};
+
+pub fn download(lua: &Lua, params: Table) -> mlua::Result<Table> {
+    let src = params.get::<String>("src")?;
+    let dst = params.get::<String>("dst")?;
+    let module = lua
+        .load(chunk! {
+            local module = komandan.KomandanModule:new({ name = "download" })
+
+            function module:run()
+                module.ssh:download($src, $dst)
+            end
+
+            return module
+        })
+        .eval::<Table>()
+        .into_lua_err()?;
+
+    Ok(module)
+}

src/modules/mod.rs

@@ -1,5 +1,9 @@
 mod cmd;
 mod script;
+mod upload;
+mod download;
 
 pub use cmd::cmd;
 pub use script::script;
+pub use upload::upload;
+pub use download::download;

src/modules/upload.rs

@@ -0,0 +1,20 @@
+use mlua::{chunk, ExternalResult, Lua, Table};
+
+pub fn upload(lua: &Lua, params: Table) -> mlua::Result<Table> {
+    let src = params.get::<String>("src")?;
+    let dst = params.get::<String>("dst")?;
+    let module = lua
+        .load(chunk! {
+            local module = komandan.KomandanModule:new({ name = "upload" })
+
+            function module:run()
+                module.ssh:upload($src, $dst)
+            end
+
+            return module
+        })
+        .eval::<Table>()
+        .into_lua_err()?;
+
+    Ok(module)
+}

src/ssh.rs

@@ -1,12 +1,13 @@
 use std::{
-    io::{Read, Write},
+    fs,
+    io::{self, Read, Write},
     net::TcpStream,
     path::Path,
 };
 
 use anyhow::Result;
 use mlua::{Lua, Table, UserData};
-use ssh2::Session;
+use ssh2::{Session, Sftp};
 
 pub struct SSHSession {
     session: Session,
@@ -80,6 +81,30 @@ impl SSHSession {
         Ok((stdout, stderr, exit_code))
     }
 
+    pub fn upload(&mut self, local_path: &Path, remote_path: &Path) -> Result<()> {
+        let mut sftp = self.session.sftp()?;
+
+        if local_path.is_dir() {
+            upload_directory(&mut sftp, local_path, remote_path)?;
+        } else {
+            upload_file(&mut sftp, local_path, remote_path)?;
+        }
+
+        Ok(())
+    }
+
+    pub fn download(&mut self, remote_path: &Path, local_path: &Path) -> Result<()> {
+        let mut sftp = self.session.sftp()?;
+
+        if remote_path.is_dir() {
+            download_directory(&mut sftp, remote_path, local_path)?;
+        } else {
+            download_file(&mut sftp, remote_path, local_path)?;
+        }
+
+        Ok(())
+    }
+
     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
@@ -96,6 +121,77 @@ impl SSHSession {
     }
 }
 
+fn upload_file(
+    sftp: &mut Sftp,
+    local_path: &Path,
+    remote_path: &Path,
+) -> io::Result<()> {
+    let mut local_file = fs::File::open(local_path)?;
+    let mut remote_file = sftp.create(remote_path)?;
+
+    io::copy(&mut local_file, &mut remote_file)?;
+
+    Ok(())
+}
+
+fn upload_directory(
+    sftp: &mut Sftp,
+    local_path: &Path,
+    remote_path: &Path,
+) -> io::Result<()> {
+    if !sftp.stat(remote_path).is_ok() {
+        sftp.mkdir(remote_path, 0o755)?;
+    }
+
+    for entry in fs::read_dir(local_path)? {
+        let entry = entry?;
+        let entry_path = entry.path();
+        let entry_name = entry.file_name();
+        let remote_entry_path = remote_path.join(entry_name);
+
+        if entry_path.is_dir() {
+            upload_directory(sftp, &entry_path, &remote_entry_path)?;
+        } else {
+            upload_file(sftp, &entry_path, &remote_entry_path)?;
+        }
+    }
+
+    Ok(())
+}
+
+fn download_file(sftp: &mut Sftp, remote_path: &Path, local_path: &Path) -> io::Result<()> {
+    let mut remote_file = sftp.open(remote_path)?;
+    let mut local_file = fs::File::create(local_path)?;
+
+    io::copy(&mut remote_file, &mut local_file)?;
+
+    Ok(())
+}
+
+fn download_directory(sftp: &mut Sftp, remote_path: &Path, local_path: &Path) -> io::Result<()> {
+    if !local_path.exists() {
+        fs::create_dir_all(local_path)?;
+    }
+
+    for entry in sftp.readdir(remote_path)? {
+        let entry_name = entry.0.file_name().unwrap();
+        let remote_entry_path = remote_path.join(entry_name);
+        let local_entry_path = local_path.join(entry_name);
+
+        if entry_name == "." || entry_name == ".." {
+            continue;
+        }
+
+        if entry.1.file_type() == ssh2::FileType::Directory {
+            download_directory(sftp, &remote_entry_path, &local_entry_path)?;
+        } else {
+            download_file(sftp, &remote_entry_path, &local_entry_path)?;
+        }
+    }
+
+    Ok(())
+}
+
 impl UserData for SSHSession {
     fn add_methods<M: mlua::UserDataMethods<Self>>(methods: &mut M) {
         methods.add_method_mut("cmd", |lua, this, command: String| {
@@ -118,6 +214,28 @@ impl UserData for SSHSession {
             },
         );
 
+        methods.add_method_mut(
+            "upload",
+            |_, this, (local_path, remote_path): (String, String)| {
+                this.upload(
+                    Path::new(local_path.as_str()),
+                    Path::new(remote_path.as_str()),
+                )?;
+                Ok(())
+            },
+        );
+
+        methods.add_method_mut(
+            "download",
+            |_, this, (remote_path, local_path): (String, String)| {
+                this.download(
+                    Path::new(remote_path.as_str()),
+                    Path::new(local_path.as_str()),
+                )?;
+                Ok(())
+            },
+        );
+
         methods.add_method("get_session_results", |lua, this, ()| {
             let table = lua.create_table()?;
             table.set("stdout", this.stdout.as_ref().unwrap().clone())?;