diff --git a/.github/workflows/publish.yaml b/.github/workflows/publish.yaml new file mode 100644 index 0000000..6ebd712 --- /dev/null +++ b/.github/workflows/publish.yaml @@ -0,0 +1,28 @@ +name: Test and Publish crate + +on: + push: + branches: [main] + pull_request: + branches: [main] + +env: + CARGO_REGISTRY_TOKEN: ${{ secrets.CARGO_REGISTRY_TOKEN }} + +jobs: + build: + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v2 + + - name: Install Rust + uses: actions-rs/toolchain@v1 + with: + toolchain: stable + + - run: cargo test + + - name: Publish + if: github.event_name == 'push' && github.ref == 'refs/heads/main' + run: cargo publish || true diff --git a/.gitignore b/.gitignore index 0b745e2..b3290a5 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,3 @@ /target -.env \ No newline at end of file +.env +temp/ \ No newline at end of file diff --git a/Cargo.lock b/Cargo.lock index 8844cb3..6c0d027 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -17,6 +17,55 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe" +[[package]] +name = "anstream" +version = "0.6.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "418c75fa768af9c03be99d17643f93f79bbba589895012a80e3452a19ddda15b" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + +[[package]] +name = "anstyle" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "038dfcf04a5feb68e9c60b21c9625a54c2c0616e79b72b0fd87075a056ae1d1b" + +[[package]] +name = "anstyle-parse" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c03a11a9034d92058ceb6ee011ce58af4a9bf61491aa7e1e59ecd24bd40d22d4" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a64c907d4e79225ac72e2a354c9ce84d50ebb4586dee56c82b3ee73004f537f5" +dependencies = [ + "windows-sys 0.52.0", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61a38449feb7068f52bb06c12759005cf459ee52bb4adc1d5a7c4322d716fb19" +dependencies = [ + "anstyle", + "windows-sys 0.52.0", +] + [[package]] name = "anyhow" version = "1.0.83" @@ -75,10 +124,204 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" [[package]] -name = "dotenvy" -version = "0.15.7" +name = "clap" +version = "4.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90bc066a67923782aa8515dbaea16946c5bcc5addbd668bb80af688e53e548a0" +dependencies = [ + "clap_builder", + "clap_derive", +] + +[[package]] +name = "clap_builder" +version = "4.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae129e2e766ae0ec03484e609954119f123cc1fe650337e155d03b022f24f7b4" +dependencies = [ + "anstream", + "anstyle", + "clap_lex", + "strsim", +] + +[[package]] +name = "clap_derive" +version = "4.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "528131438037fd55894f62d6e9f068b8f45ac57ffa77517819645d10aed04f64" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "clap_lex" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "98cc8fbded0c607b7ba9dd60cd98df59af97e84d24e49c8557331cfc26d301ce" + +[[package]] +name = "colorchoice" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b6a852b24ab71dffc585bcb46eaf7959d175cb865a7152e35b348d1b2960422" + +[[package]] +name = "console" +version = "0.15.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e1f83fc076bd6dd27517eacdf25fef6c4dfe5f1d7448bafaaf3a26f13b5e4eb" +dependencies = [ + "encode_unicode", + "lazy_static", + "libc", + "unicode-width", + "windows-sys 0.52.0", +] + +[[package]] +name = "dirs-next" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b98cf8ebf19c3d1b223e151f99a4f9f0690dca41414773390fc824184ac833e1" +dependencies = [ + "cfg-if", + "dirs-sys-next", +] + +[[package]] +name = "dirs-sys-next" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ebda144c4fe02d1f7ea1a7d9641b6fc6b580adcfa024ae48797ecdeb6825b4d" +dependencies = [ + "libc", + "redox_users", + "winapi", +] + +[[package]] +name = "encode_unicode" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a357d28ed41a50f9c765dbfe56cbc04a64e53e5fc58ba79fbc34c10ef3df831f" + +[[package]] +name = "fast-scp" +version = "0.1.0-beta.1" +dependencies = [ + "anyhow", + "clap", + "dirs-next", + "futures", + "indicatif", + "insta", + "ssh2", + "tokio", +] + +[[package]] +name = "futures" +version = "0.3.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "645c6916888f6cb6350d2550b80fb63e734897a8498abe35cfb732b6487804b0" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-channel" +version = "0.3.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eac8f7d7865dcb88bd4373ab671c8cf4508703796caa2b1985a9ca867b3fcb78" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfc6580bb841c5a68e9ef15c77ccc837b40a7504914d52e47b8b0e9bbda25a1d" + +[[package]] +name = "futures-executor" +version = "0.3.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a576fc72ae164fca6b9db127eaa9a9dda0d61316034f33a0a0d4eda41f02b01d" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a44623e20b9681a318efdd71c299b6b222ed6f231972bfe2f224ebad6311f0c1" + +[[package]] +name = "futures-macro" +version = "0.3.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87750cf4b7a4c0625b1529e4c543c2182106e4dedc60a2a6455e00d212c489ac" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "futures-sink" +version = "0.3.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fb8e00e87438d937621c1c6269e53f536c14d3fbd6a042bb24879e57d474fb5" + +[[package]] +name = "futures-task" +version = "0.3.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38d84fa142264698cdce1a9f9172cf383a0c82de1bddcf3092901442c4097004" + +[[package]] +name = "futures-util" +version = "0.3.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d6401deb83407ab3da39eba7e33987a73c3df0c82b4bb5813ee871c19c41d48" +dependencies = [ + "futures-channel", + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "pin-utils", + "slab", +] + +[[package]] +name = "getrandom" +version = "0.2.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1aaf95b3e5c8f23aa320147307562d361db0ae0d51242340f558153b4eb2439b" +checksum = "c4567c8db10ae91089c99af84c68c38da3ec2f087c3f82960bcdbf3656b6f4d7" +dependencies = [ + "cfg-if", + "libc", + "wasi", +] [[package]] name = "gimli" @@ -86,12 +329,43 @@ version = "0.28.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4271d37baee1b8c7e4b708028c57d816cf9d2434acb33a549475f78c181f6253" +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + [[package]] name = "hermit-abi" version = "0.3.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d231dfb89cfffdbc30e7fc41579ed6066ad03abda9e567ccafae602b97ec5024" +[[package]] +name = "indicatif" +version = "0.17.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "763a5a8f45087d6bcea4222e7b72c291a054edf80e4ef6efd2a4979878c7bea3" +dependencies = [ + "console", + "instant", + "number_prefix", + "portable-atomic", + "unicode-width", +] + +[[package]] +name = "insta" +version = "1.38.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3eab73f58e59ca6526037208f0e98851159ec1633cf17b6cd2e1f2c3fd5d53cc" +dependencies = [ + "console", + "lazy_static", + "linked-hash-map", + "similar", +] + [[package]] name = "instant" version = "0.1.12" @@ -101,12 +375,34 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "is_terminal_polyfill" +version = "1.70.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8478577c03552c21db0e2724ffb8986a5ce7af88107e6be5d2ee6e158c12800" + +[[package]] +name = "lazy_static" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2abad23fbc42b3700f2f279844dc832adb2b2eb069b2df918f455c4e18cc646" + [[package]] name = "libc" version = "0.2.154" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ae743338b92ff9146ce83992f766a31066a91a8c84a45e0e9f21e7cf6de6d346" +[[package]] +name = "libredox" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0ff37bd590ca25063e35af745c343cb7a0271906fb7b37e4813e8f79f00268d" +dependencies = [ + "bitflags 2.5.0", + "libc", +] + [[package]] name = "libssh2-sys" version = "0.3.0" @@ -133,6 +429,12 @@ dependencies = [ "vcpkg", ] +[[package]] +name = "linked-hash-map" +version = "0.5.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0717cef1bc8b636c6e1c1bbdefc09e6322da8a9321966e8928ef80d20f7f770f" + [[package]] name = "lock_api" version = "0.4.12" @@ -179,6 +481,12 @@ dependencies = [ "libc", ] +[[package]] +name = "number_prefix" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3" + [[package]] name = "object" version = "0.32.2" @@ -254,12 +562,24 @@ version = "0.2.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bda66fc9667c18cb2758a2ac84d1167245054bcf85d5d1aaa6923f45801bdd02" +[[package]] +name = "pin-utils" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" + [[package]] name = "pkg-config" version = "0.3.30" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d231b230927b5e4ad203db57bbcbee2802f6bce620b1e4a9024a07d94e2907ec" +[[package]] +name = "portable-atomic" +version = "1.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7170ef9988bc169ba16dd36a7fa041e5c4cbeb6a35b76d4c03daded371eae7c0" + [[package]] name = "proc-macro2" version = "1.0.82" @@ -296,6 +616,17 @@ dependencies = [ "bitflags 2.5.0", ] +[[package]] +name = "redox_users" +version = "0.4.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bd283d9651eeda4b2a83a43c1c91b266c40fd76ecd39a50a8c630ae69dc72891" +dependencies = [ + "getrandom", + "libredox", + "thiserror", +] + [[package]] name = "rustc-demangle" version = "0.1.24" @@ -308,16 +639,6 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" -[[package]] -name = "scp-cli" -version = "0.1.0" -dependencies = [ - "anyhow", - "dotenvy", - "ssh2", - "tokio", -] - [[package]] name = "signal-hook-registry" version = "1.4.2" @@ -327,6 +648,21 @@ dependencies = [ "libc", ] +[[package]] +name = "similar" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fa42c91313f1d05da9b26f267f931cf178d4aba455b4c4622dd7355eb80c6640" + +[[package]] +name = "slab" +version = "0.4.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f92a496fb766b417c996b9c5e57daf2f7ad3b0bebe1ccfca4856390e3d3bb67" +dependencies = [ + "autocfg", +] + [[package]] name = "smallvec" version = "1.13.2" @@ -355,6 +691,12 @@ dependencies = [ "parking_lot 0.11.2", ] +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + [[package]] name = "syn" version = "2.0.61" @@ -366,6 +708,26 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "thiserror" +version = "1.0.60" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "579e9083ca58dd9dcf91a9923bb9054071b9ebbd800b342194c9feb0ee89fc18" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.60" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2470041c06ec3ac1ab38d0356a6119054dedaea53e12fbefc0de730a1c08524" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "tokio" version = "1.37.0" @@ -402,6 +764,18 @@ version = "1.0.12" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3354b9ac3fae1ff6755cb6db53683adb661634f67557942dea4facebec0fee4b" +[[package]] +name = "unicode-width" +version = "0.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68f5e5f3158ecfd4b8ff6fe086db7c8467a2dfdac97fe420f2b7c4aa97af66d6" + +[[package]] +name = "utf8parse" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "711b9620af191e0cdc7468a8d14e709c3dcdb115b36f838e601583af800a370a" + [[package]] name = "vcpkg" version = "0.2.15" diff --git a/Cargo.toml b/Cargo.toml index f8fff49..eb1bc4e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,10 +1,14 @@ [package] -name = "scp-cli" -version = "0.1.0" +name = "fast-scp" +version = "0.1.0-beta.1" edition = "2021" [dependencies] ssh2 = "0.9.4" tokio = { version = "1.37.0", features = ["full"] } anyhow = "1.0.83" -dotenvy = "0.15.7" +clap = { version = "4.5.4", features = ["derive"] } +dirs-next = "2.0.0" +indicatif = "0.17.8" +futures = "0.3.30" +insta = "1.38.0" diff --git a/README.md b/README.md new file mode 100644 index 0000000..c2605a2 --- /dev/null +++ b/README.md @@ -0,0 +1,11 @@ +# scp-rs + +**WARNING: THIS CLI TOOL IS STILL IN BETA AND NOT READY FOR USE** + +A Rust CLI tool to copy files from remote server to local machine or the other way around, handles tasks concurrently, which makes it faster than the traditional scp command. + +## Example + +```bash +fast-scp receive --host --user --private-key [path-to-private-key] +``` diff --git a/src/cli.rs b/src/cli.rs new file mode 100644 index 0000000..69c1d8c --- /dev/null +++ b/src/cli.rs @@ -0,0 +1,33 @@ +use clap::{command, Parser, Subcommand}; +use std::path::PathBuf; + +#[derive(Parser)] +#[command(version, about, long_about = None)] +pub struct Cli { + #[command(subcommand)] + pub command: Commands, +} + +#[derive(Subcommand)] +pub enum Commands { + #[command(about = "Copy a file to a remote host")] + Receive { + #[clap(help = "Remote source to copy from")] + source: String, + + #[clap(help = "Local destination to copy to")] + destination: String, + + #[clap(long, help = "Remote host to connect to")] + host: String, + + #[clap(short, long, help = "Remote username to connect as")] + user: String, + + #[clap(short, long, help = "Path to private key")] + private_key: Option, + + #[clap(help = "Replace the file if it exists", long, default_value = "false")] + replace: bool, + }, +} diff --git a/src/error.rs b/src/error.rs new file mode 100644 index 0000000..c12e6f0 --- /dev/null +++ b/src/error.rs @@ -0,0 +1,31 @@ +use core::fmt; +use std::fmt::{Display, Formatter}; + +#[derive(Debug)] +pub enum ScpError { + Io(std::io::Error), + Ssh(ssh2::Error), + Other(String), +} + +impl Display for ScpError { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + match self { + ScpError::Io(e) => write!(f, "IO error: {}", e), + ScpError::Ssh(e) => write!(f, "SSH error: {}", e), + ScpError::Other(e) => write!(f, "Error: {}", e), + } + } +} + +impl From for ScpError { + fn from(e: std::io::Error) -> Self { + ScpError::Io(e) + } +} + +impl From for ScpError { + fn from(e: ssh2::Error) -> Self { + ScpError::Ssh(e) + } +} diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..bd1733c --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,5 @@ +pub mod cli; +pub mod error; +pub mod run; +pub mod scp; +pub mod utils; diff --git a/src/main.rs b/src/main.rs index 2963f1b..cc1bf01 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,51 +1,9 @@ -use dotenvy::dotenv; -use ssh2::Session; -use std::fs::File; -use std::io::prelude::*; -use std::net::TcpStream; -use std::path::Path; - -fn copy_file_from_vps( - host: &str, - username: &str, - private_key_path: &str, - remote_file_path: &str, - local_file_path: &str, -) -> anyhow::Result<()> { - // Connect to the host - let tcp = TcpStream::connect(host)?; - let mut session = Session::new()?; - session.set_tcp_stream(tcp); - session.handshake()?; - - // Authenticate using a private key - session.userauth_pubkey_file(username, None, Path::new(private_key_path), None)?; - - // Create a SCP channel for receiving the file - let (mut remote_file, stat) = session.scp_recv(Path::new(remote_file_path))?; - let mut contents = Vec::with_capacity(stat.size() as usize); - remote_file.read_to_end(&mut contents)?; - - // Create local file and write to it - let mut local_file = File::create(Path::new(local_file_path))?; - local_file.write_all(&contents)?; - - Ok(()) -} +use fast_scp::run::run; #[tokio::main] -async fn main() -> anyhow::Result<()> { - dotenv().ok(); - - let host = std::env::var("VPS_HOST")?; - let username = std::env::var("VPS_USERNAME")?; - let private_key_path = std::env::var("VPS_PRIVATE_KEY_PATH")?; - - copy_file_from_vps( - &format!("{}:22", host), - &username, - &private_key_path, - "/path/to/remote/file", - "/path/to/local/file", - ) +async fn main() { + match run().await { + Ok(_) => (), + Err(e) => eprintln!("Error: {}", e), + } } diff --git a/src/run.rs b/src/run.rs new file mode 100644 index 0000000..86b451b --- /dev/null +++ b/src/run.rs @@ -0,0 +1,35 @@ +use crate::cli::{Cli, Commands}; +use crate::error::ScpError; +use crate::scp::{Connect, Mode, SshOpts}; +use crate::utils::get_private_key_path; +use clap::Parser; +use std::path::PathBuf; + +pub async fn run() -> anyhow::Result<(), ScpError> { + let args = Cli::parse(); + + match args.command { + Commands::Receive { + source, + destination, + host, + user: username, + private_key, + replace, + } => { + let private_key = get_private_key_path(&private_key)?; + + let scp_opts = SshOpts { + host: format!("{}:22", host), + private_key, + username, + }; + + let mode = if replace { Mode::Replace } else { Mode::Ignore }; + + return Connect::new(scp_opts, mode)? + .receive(&PathBuf::from(source), &PathBuf::from(destination)) + .await; + } + } +} diff --git a/src/scp.rs b/src/scp.rs new file mode 100644 index 0000000..0f6bfd5 --- /dev/null +++ b/src/scp.rs @@ -0,0 +1,181 @@ +use futures::future::join_all; +use indicatif::ProgressBar; +use ssh2::Session; +use std::{ + fs::{self, File}, + io::{Read, Write}, + net::TcpStream, + path::PathBuf, +}; + +use crate::{error::ScpError, utils::with_retry}; + +pub struct Connect { + session: Session, + ssh_opts: SshOpts, + mode: Mode, +} + +impl Connect { + pub fn new(ssh_opts: SshOpts, mode: Mode) -> anyhow::Result { + let session = create_session(&ssh_opts)?; + + Ok(Self { + session, + ssh_opts, + mode, + }) + } + + pub async fn receive(&self, from: &PathBuf, to: &PathBuf) -> anyhow::Result<(), ScpError> { + let start = std::time::Instant::now(); + + let files = self.list(from)?; + let pb = ProgressBar::new(files.len() as u64); + + let mut handles = Vec::new(); + for item in files { + let to_path = to.join(item.strip_prefix(from).unwrap()); + let item_clone = item.clone(); + let ssh_opts = self.ssh_opts.clone(); + let pb = pb.clone(); + let mode = self.mode.clone(); + let handle = tokio::task::spawn(async move { + let result = + copy_file_from_remote(&ssh_opts, item_clone.clone(), to_path, &mode).await; + pb.inc(1); + result + }); + + handles.push(handle); + } + + let items = join_all(handles).await; + + if items.iter().all(|x| x.is_ok()) { + println!("Done in {:.2?}", start.elapsed()); + Ok(()) + } else { + Err(std::io::Error::new( + std::io::ErrorKind::Other, + "One or more files failed to copy", + ) + .into()) + } + } + + fn list(&self, dir: &PathBuf) -> anyhow::Result, ScpError> { + let mut channel = self.session.channel_session()?; + + channel.exec(&format!("ls -R {}", dir.display()))?; + + let mut buf = String::new(); + channel.read_to_string(&mut buf)?; + + let files_only = find_files(&buf); + + Ok(files_only) + } +} + +pub fn find_files(buf: &str) -> Vec { + let mut dirs: Vec = Vec::new(); + let structured = buf + .split("\n\n") + .map(|x| { + let mut lines = x.lines(); + let dir: PathBuf = lines.next().unwrap().split(":").next().unwrap().into(); + + let files = lines.collect::>(); + + let full_path = files + .iter() + .map(|x| PathBuf::new().join(x)) + .map(|x| dir.join(x)) + .collect::>(); + + dirs.push(dir); + full_path + }) + .collect::>(); + + let flattened = structured.iter().flatten().collect::>(); + + let files_only = flattened + .iter() + .filter(|x| !dirs.contains(x)) + .map(|x| x.to_path_buf()) + .collect::>(); + + files_only +} + +#[derive(Clone)] +pub struct SshOpts { + pub host: String, + pub username: String, + pub private_key: PathBuf, +} + +/// Mode to use when copying files +/// Replace will overwrite the file if it exists +/// Ignore will skip the file if it exists +#[derive(Clone)] +pub enum Mode { + Replace, + Ignore, +} + +async fn copy_file_from_remote( + ssh_opts: &SshOpts, + remote_file_path: PathBuf, + local_file_path: PathBuf, + mode: &Mode, +) -> anyhow::Result<(), ScpError> { + let create_session = || create_session(ssh_opts); + let session = with_retry(create_session, 10)?; + + // Create a SCP channel for receiving the file + let (mut remote_file, stat) = session.scp_recv(&remote_file_path)?; + let mut contents = Vec::with_capacity(stat.size() as usize); + remote_file.read_to_end(&mut contents)?; + + // make the dir if not exists + fs::create_dir_all(local_file_path.parent().unwrap())?; + + match mode { + Mode::Replace => { + let mut local_file = File::create(&local_file_path)?; + local_file.write_all(&contents)?; + } + Mode::Ignore => { + if local_file_path.exists() { + println!( + "Skipping already existing file: {}", + local_file_path.display() + ); + return Ok(()); + } + + let mut local_file = File::create(local_file_path)?; + local_file.write_all(&contents)?; + } + } + + session.disconnect(None, "Bye", None)?; + + Ok(()) +} + +pub fn create_session(ssh_opts: &SshOpts) -> anyhow::Result { + // Connect to the host + let tcp = TcpStream::connect(&ssh_opts.host)?; + let mut session = Session::new()?; + session.set_tcp_stream(tcp); + session.handshake()?; + + // Authenticate using a private key + session.userauth_pubkey_file(&ssh_opts.username, None, &ssh_opts.private_key, None)?; + + Ok(session) +} diff --git a/src/utils.rs b/src/utils.rs new file mode 100644 index 0000000..ba0598a --- /dev/null +++ b/src/utils.rs @@ -0,0 +1,39 @@ +use dirs_next::home_dir; +use std::path::PathBuf; + +use crate::error::ScpError; + +pub fn with_retry(f: F, max_retries: u32) -> anyhow::Result +where + F: Fn() -> anyhow::Result, +{ + let mut retries = 0; + loop { + match f() { + Ok(x) => return Ok(x), + Err(e) => { + if retries >= max_retries { + return Err(e); + } + + retries += 1; + } + } + } +} + +pub fn get_private_key_path(private_key: &Option) -> anyhow::Result { + match private_key { + Some(path) => Ok(PathBuf::from(path)), + None => Ok(home_dir() + .ok_or( + ScpError::Io( + std::io::Error::new( + std::io::ErrorKind::Other, + "Could not find home directory, please provide the private key path using the --private-key-path flag", + ), + ), + )? + .join(".ssh/id_rsa")), + } +} diff --git a/tests/files.rs b/tests/files.rs new file mode 100644 index 0000000..0d5dc6a --- /dev/null +++ b/tests/files.rs @@ -0,0 +1,20 @@ +use std::fs; + +fn open_file() -> String { + fs::read_to_string("tests/snapshots/ls-R.log").expect("Unable to read file") +} + +#[cfg(test)] +mod tests { + use super::*; + use fast_scp::scp::*; + use insta; + + #[test] + fn test_find_files() { + let file = open_file(); + let files_only = find_files(&file); + println!("{:?}", files_only); + insta::assert_debug_snapshot!(files_only); + } +} diff --git a/tests/snapshots/files__tests__find_files.snap b/tests/snapshots/files__tests__find_files.snap new file mode 100644 index 0000000..677accb --- /dev/null +++ b/tests/snapshots/files__tests__find_files.snap @@ -0,0 +1,25 @@ +--- +source: tests/files.rs +assertion_line: 18 +expression: files_only +--- +[ + "project/README.md", + "project/package.json", + "project/src/app/App.tsx", + "project/src/app/index.ts", + "project/src/components/Header.tsx", + "project/src/components/Footer.tsx", + "project/src/components/Sidebar.tsx", + "project/src/components/index.ts", + "project/src/hooks/useAuth.ts", + "project/src/hooks/useFetch.ts", + "project/src/hooks/index.ts", + "project/src/utils/helpers.ts", + "project/src/utils/constants.ts", + "project/src/utils/index.ts", + "project/tests/App.test.tsx", + "project/tests/Header.test.tsx", + "project/tests/Footer.test.tsx", + "project/tests/Sidebar.test.tsx", +] diff --git a/tests/snapshots/ls-R.log b/tests/snapshots/ls-R.log new file mode 100644 index 0000000..128d8d8 --- /dev/null +++ b/tests/snapshots/ls-R.log @@ -0,0 +1,37 @@ +project: +src +tests +README.md +package.json + +project/src: +app +components +hooks +utils + +project/src/app: +App.tsx +index.ts + +project/src/components: +Header.tsx +Footer.tsx +Sidebar.tsx +index.ts + +project/src/hooks: +useAuth.ts +useFetch.ts +index.ts + +project/src/utils: +helpers.ts +constants.ts +index.ts + +project/tests: +App.test.tsx +Header.test.tsx +Footer.test.tsx +Sidebar.test.tsx \ No newline at end of file