diff --git a/Cargo.toml b/Cargo.toml index f83b8e7..17a0b26 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -18,7 +18,7 @@ pyo3 = { version = "0.26.0", features = ["generate-import-lib"] } shellexpand = "3.1.0" pyo3-async-runtimes = { version = "0.26.0", features = ["tokio-runtime"] } tokio = { version = "1", features = ["full"] } -russh = { version = "0.56", default-features = false, features = ["flate2", "async-trait", "ring", "rsa"] } +russh = { version = "0.58", default-features = false, features = ["flate2", "async-trait", "ring", "rsa"] } russh-sftp = "2.1" async-trait = "0.1" diff --git a/README.md b/README.md index ed8804b..0d71e32 100644 --- a/README.md +++ b/README.md @@ -85,6 +85,12 @@ conn.sftp_write(local_path="/path/to/my/file", remote_path="/dest/path/file") # Read a remote file contents = conn.sftp_read(remote_path="/dest/path/file") + +# Upload an entire local directory recursively +transferred, failed = conn.sftp_put_dir("/local/build/", "/remote/app/") + +# Download an entire remote directory recursively +transferred, failed = conn.sftp_get_dir("/remote/logs/", "/local/logs/") ``` 📚 **For complete documentation including SCP, file tailing, interactive shells, and more, see [Synchronous Usage](docs/synchronous.md).** diff --git a/docs/asynchronous.md b/docs/asynchronous.md index 03921fd..ea08518 100644 --- a/docs/asynchronous.md +++ b/docs/asynchronous.md @@ -207,6 +207,35 @@ async with AsyncConnection(host="my.test.server", password="pass") as conn: print(file) ``` +### Directory Transfers + +Upload or download entire directory trees with a single awaitable call. Both methods return a tuple of `(transferred_files, failed_files)`, where each is a list of file paths. + +```python +async with AsyncConnection(host="my.test.server", password="pass") as conn: + # Upload an entire local directory to the remote server + transferred, failed = await conn.sftp_put_dir( + local_path="/local/build/", + remote_path="/remote/app/", + ) + if failed: + print(f"Failed to upload: {failed}") + + # Download an entire remote directory to a local destination + transferred, failed = await conn.sftp_get_dir( + remote_path="/remote/logs/", + local_path="/local/logs/", + ) +``` + +Optional keyword arguments control symlink, permission, and error-handling behaviour: + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `follow_symlinks` | `True` | Follow symlinks; set `False` to skip them | +| `preserve_permissions` | `True` | Mirror source permissions on the destination | +| `fail_fast` | `False` | Raise on the first error instead of collecting failures | + ### Concurrent File Operations ```python diff --git a/docs/synchronous.md b/docs/synchronous.md index 1e05b44..a2b0b29 100644 --- a/docs/synchronous.md +++ b/docs/synchronous.md @@ -123,6 +123,42 @@ contents = conn.sftp_read(remote_path="/dest/path/file") print(contents) ``` +### Directory Transfers + +Upload or download entire directory trees with a single call. Both methods return a tuple of `(transferred_files, failed_files)`, where each is a list of file paths. + +```python +# Upload an entire local directory to the remote server +transferred, failed = conn.sftp_put_dir( + local_path="/local/build/", + remote_path="/remote/app/", +) +if failed: + print(f"Failed to upload: {failed}") + +# Download an entire remote directory to a local destination +transferred, failed = conn.sftp_get_dir( + remote_path="/remote/logs/", + local_path="/local/logs/", +) +``` + +Optional keyword arguments control symlink, permission, and error-handling behaviour: + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `follow_symlinks` | `True` | Follow symlinks; set `False` to skip them | +| `preserve_permissions` | `True` | Mirror source permissions on the destination | +| `fail_fast` | `False` | Raise on the first error instead of collecting failures | + +```python +# Raise immediately on the first error +conn.sftp_put_dir("/src/", "/dst/", fail_fast=True) + +# Skip symlinks and don't mirror permissions +conn.sftp_put_dir("/src/", "/dst/", follow_symlinks=False, preserve_permissions=False) +``` + ### Copy Files Between Connections Hussh provides a convenient way to copy files between two remote servers: diff --git a/src/asynchronous.rs b/src/asynchronous.rs index 808a97e..991fbe9 100644 --- a/src/asynchronous.rs +++ b/src/asynchronous.rs @@ -79,6 +79,12 @@ //! //! # List directory contents //! files = await conn.sftp_list("/remote/path") +//! +//! # Upload an entire local directory recursively +//! transferred, failed = await conn.sftp_put_dir("/local/build/", "/remote/app/") +//! +//! # Download an entire remote directory recursively +//! transferred, failed = await conn.sftp_get_dir("/remote/logs/", "/local/logs/") //! ``` //! //! For file tailing: @@ -92,7 +98,7 @@ //! print(tailer.contents) //! ``` -use crate::connection::SSHResult; +use crate::connection::{SSHResult, MAX_BUFF_SIZE}; use pyo3::exceptions::{PyRuntimeError, PyTimeoutError}; use pyo3::prelude::*; use russh::client::{Config, Handle, Handler}; @@ -222,6 +228,26 @@ impl Handler for ClientHandler { /// /// * `path`: The path to the remote directory to list. /// +/// ### `sftp_put_dir` +/// +/// Uploads a local directory recursively to a remote path over SFTP. Returns (transferred_files, failed_files). It takes the following parameters: +/// +/// * `local_path`: The local directory to upload. +/// * `remote_path`: The remote destination path. +/// * `follow_symlinks`: Whether to follow symlinks (default: true). +/// * `preserve_permissions`: Whether to preserve file permissions (default: true). +/// * `fail_fast`: Whether to raise an exception on the first error instead of collecting failures (default: false). +/// +/// ### `sftp_get_dir` +/// +/// Downloads a remote directory recursively to a local path over SFTP. Returns (transferred_files, failed_files). It takes the following parameters: +/// +/// * `remote_path`: The remote directory to download. +/// * `local_path`: The local destination path. +/// * `follow_symlinks`: Whether to follow symlinks (default: true). +/// * `preserve_permissions`: Whether to preserve file permissions (default: true). +/// * `fail_fast`: Whether to raise an exception on the first error instead of collecting failures (default: false). +/// /// ### `shell` /// /// Creates an `AsyncInteractiveShell` instance. It takes the following parameter: @@ -492,7 +518,7 @@ impl AsyncConnection { PyRuntimeError::new_err(format!("Failed to create local file: {}", e)) })?; - let mut buffer = vec![0u8; 65536]; // 64KB buffer to match sync version + let mut buffer = vec![0u8; MAX_BUFF_SIZE]; loop { let n = remote_file.read(&mut buffer).await.map_err(|e| { PyRuntimeError::new_err(format!("Failed to read remote file: {}", e)) @@ -536,7 +562,7 @@ impl AsyncConnection { .await .map_err(|e| PyRuntimeError::new_err(format!("Failed to create remote file: {}", e)))?; - let mut buffer = vec![0u8; 65536]; // 64KB buffer to match sync version + let mut buffer = vec![0u8; MAX_BUFF_SIZE]; loop { let n = local_file.read(&mut buffer).await.map_err(|e| { PyRuntimeError::new_err(format!("Failed to read local file: {}", e)) @@ -548,6 +574,14 @@ impl AsyncConnection { PyRuntimeError::new_err(format!("Failed to write remote file: {}", e)) })? } + remote_file + .flush() + .await + .map_err(|e| PyRuntimeError::new_err(format!("Failed to flush remote file: {}", e)))?; + remote_file + .shutdown() + .await + .map_err(|e| PyRuntimeError::new_err(format!("Failed to close remote file: {}", e)))?; Ok(()) } @@ -571,10 +605,393 @@ impl AsyncConnection { .await .map_err(|e| PyRuntimeError::new_err(format!("Failed to write remote file: {}", e)))?; + remote_file + .flush() + .await + .map_err(|e| PyRuntimeError::new_err(format!("Failed to flush remote file: {}", e)))?; + remote_file + .shutdown() + .await + .map_err(|e| PyRuntimeError::new_err(format!("Failed to close remote file: {}", e)))?; + Ok(()) } - /// Create a new AsyncConnection - used by MultiConnection + /// Internal async SFTP put_dir implementation + pub(crate) async fn sftp_put_dir_async( + &self, + local_path: String, + remote_path: String, + follow_symlinks: bool, + preserve_permissions: bool, + fail_fast: bool, + ) -> PyResult<(Vec, Vec)> { + let sftp = + AsyncConnection::get_or_init_sftp(self.session.clone(), self.sftp_session.clone()) + .await?; + + let mut transferred: Vec = Vec::new(); + let mut failed: Vec = Vec::new(); + + // Ensure remote base directory (and all missing parents) exist — like mkdir -p + { + let parts: Vec<&str> = remote_path.split('/').filter(|s| !s.is_empty()).collect(); + let is_absolute = remote_path.starts_with('/'); + let mut current = if is_absolute { + String::from("/") + } else { + String::new() + }; + for part in &parts { + if !current.is_empty() && !current.ends_with('/') { + current.push('/'); + } + current.push_str(part); + if !sftp.try_exists(¤t).await.unwrap_or(false) { + if let Err(e) = sftp.create_dir(¤t).await { + let msg = format!("Failed to create remote directory '{}': {}", current, e); + if fail_fast { + return Err(PyRuntimeError::new_err(msg)); + } + failed.push(local_path.clone()); + return Ok((transferred, failed)); + } + } + } + } + + // Stack for depth-first traversal: (local_dir, remote_dir_str). + // Remote paths are kept as POSIX strings to avoid OS-native path separators on Windows. + let mut dirs_to_process: Vec<(std::path::PathBuf, String)> = + vec![(std::path::PathBuf::from(&local_path), remote_path.clone())]; + + while let Some((local_dir, remote_dir)) = dirs_to_process.pop() { + let mut read_dir = match tokio::fs::read_dir(&local_dir).await { + Ok(d) => d, + Err(e) => { + let msg = format!( + "Failed to read local directory '{}': {}", + local_dir.display(), + e + ); + if fail_fast { + return Err(PyRuntimeError::new_err(msg)); + } + failed.push(local_dir.to_string_lossy().to_string()); + continue; + } + }; + + loop { + let entry = match read_dir.next_entry().await { + Ok(Some(e)) => e, + Ok(None) => break, + Err(e) => { + if fail_fast { + return Err(PyRuntimeError::new_err(format!( + "Directory entry error: {}", + e + ))); + } + continue; + } + }; + + let local_entry = entry.path(); + let local_entry_str = local_entry.to_string_lossy().to_string(); + let file_name_str = entry.file_name().to_string_lossy().to_string(); + let remote_entry_str = format!("{}/{}", remote_dir, file_name_str); + + let metadata = match if follow_symlinks { + tokio::fs::metadata(&local_entry).await + } else { + tokio::fs::symlink_metadata(&local_entry).await + } { + Ok(m) => m, + Err(e) => { + if fail_fast { + return Err(PyRuntimeError::new_err(format!( + "Metadata error for '{}': {}", + local_entry_str, e + ))); + } + failed.push(local_entry_str); + continue; + } + }; + + if metadata.is_dir() { + if !sftp.try_exists(&remote_entry_str).await.unwrap_or(false) { + match sftp.create_dir(&remote_entry_str).await { + Ok(_) => {} + Err(e) => { + if fail_fast { + return Err(PyRuntimeError::new_err(format!( + "Failed to create remote directory '{}': {}", + remote_entry_str, e + ))); + } + failed.push(local_entry_str); + continue; + } + } + } + #[cfg(unix)] + if preserve_permissions { + use std::os::unix::fs::PermissionsExt; + let mode = metadata.permissions().mode(); + // Use explicit None for all fields except permissions. + // FileAttributes::default() sets size: Some(0) which would + // truncate the target via SSH_FXP_SETSTAT. + let attrs = russh_sftp::client::fs::Metadata { + size: None, + uid: None, + user: None, + gid: None, + group: None, + permissions: Some(mode), + atime: None, + mtime: None, + }; + let _ = sftp.set_metadata(&remote_entry_str, attrs).await; + } + dirs_to_process.push((local_entry, remote_entry_str)); + } else if metadata.is_file() { + let transfer_result: Result<(), String> = async { + let mut local_file = tokio::fs::File::open(&local_entry) + .await + .map_err(|e| format!("File open error: {}", e))?; + let mut remote_file = sftp + .create(&remote_entry_str) + .await + .map_err(|e| format!("Remote file creation error: {}", e))?; + let mut buffer = vec![0u8; MAX_BUFF_SIZE]; + loop { + let n = local_file + .read(&mut buffer) + .await + .map_err(|e| format!("File read error: {}", e))?; + if n == 0 { + break; + } + remote_file + .write_all(&buffer[..n]) + .await + .map_err(|e| format!("Remote write error: {}", e))?; + } + // Flush any buffered data then close the SFTP file handle. + // shutdown() sets closed=true so Drop won't send a redundant close. + remote_file + .flush() + .await + .map_err(|e| format!("Remote file flush error: {}", e))?; + remote_file + .shutdown() + .await + .map_err(|e| format!("Remote file close error: {}", e))?; + #[cfg(unix)] + if preserve_permissions { + use std::os::unix::fs::PermissionsExt; + let mode = metadata.permissions().mode(); + // Use explicit None for all fields except permissions. + // FileAttributes::default() sets size: Some(0) which would + // truncate the file via SSH_FXP_SETSTAT. + let attrs = russh_sftp::client::fs::Metadata { + size: None, + uid: None, + user: None, + gid: None, + group: None, + permissions: Some(mode), + atime: None, + mtime: None, + }; + let _ = sftp.set_metadata(&remote_entry_str, attrs).await; + } + Ok(()) + } + .await; + match transfer_result { + Ok(_) => transferred.push(local_entry_str), + Err(e) => { + if fail_fast { + return Err(PyRuntimeError::new_err(e)); + } + failed.push(local_entry_str); + } + } + } + // Symlinks with follow_symlinks=false are skipped + } + } + + Ok((transferred, failed)) + } + + /// Internal async SFTP get_dir implementation + pub(crate) async fn sftp_get_dir_async( + &self, + remote_path: String, + local_path: String, + follow_symlinks: bool, + preserve_permissions: bool, + fail_fast: bool, + ) -> PyResult<(Vec, Vec)> { + let sftp = + AsyncConnection::get_or_init_sftp(self.session.clone(), self.sftp_session.clone()) + .await?; + + let mut transferred: Vec = Vec::new(); + let mut failed: Vec = Vec::new(); + + // Ensure local base directory exists + match tokio::fs::create_dir_all(&local_path).await { + Ok(_) => {} + Err(e) => { + let msg = format!("Failed to create local directory '{}': {}", local_path, e); + if fail_fast { + return Err(PyRuntimeError::new_err(msg)); + } + failed.push(remote_path); + return Ok((transferred, failed)); + } + } + + // Stack for depth-first traversal: (remote_dir_str, local_dir). + // Remote paths are kept as POSIX strings to avoid OS-native path separators on Windows. + let mut dirs_to_process: Vec<(String, std::path::PathBuf)> = + vec![(remote_path.clone(), std::path::PathBuf::from(&local_path))]; + + while let Some((remote_dir, local_dir)) = dirs_to_process.pop() { + let remote_dir_str = remote_dir.clone(); + let read_dir = match sftp.read_dir(&remote_dir_str).await { + Ok(d) => d, + Err(e) => { + let msg = format!( + "Failed to read remote directory '{}': {}", + remote_dir_str, e + ); + if fail_fast { + return Err(PyRuntimeError::new_err(msg)); + } + failed.push(remote_dir_str); + continue; + } + }; + + for entry in read_dir { + let file_name = entry.file_name(); + // Build remote path with POSIX separator to stay cross-platform. + let remote_entry_str = format!("{}/{}", remote_dir, file_name); + let local_entry = local_dir.join(&file_name); + + // When follow_symlinks=true, resolve symlinks on remote via metadata() + let file_type = if follow_symlinks && entry.file_type().is_symlink() { + sftp.metadata(&remote_entry_str) + .await + .map(|m| m.file_type()) + .unwrap_or_else(|_| entry.file_type()) + } else { + entry.file_type() + }; + + if file_type.is_symlink() { + // follow_symlinks=false: skip symlinks + continue; + } else if file_type.is_dir() { + match tokio::fs::create_dir_all(&local_entry).await { + Ok(_) => {} + Err(e) => { + if fail_fast { + return Err(PyRuntimeError::new_err(format!( + "Failed to create local directory '{}': {}", + local_entry.display(), + e + ))); + } + failed.push(remote_entry_str); + continue; + } + } + #[cfg(unix)] + if preserve_permissions { + if let Some(perm) = sftp + .metadata(&remote_entry_str) + .await + .ok() + .and_then(|m| m.permissions) + { + use std::os::unix::fs::PermissionsExt; + let _ = tokio::fs::set_permissions( + &local_entry, + std::fs::Permissions::from_mode(perm & 0o7777), + ) + .await; + } + } + dirs_to_process.push((remote_entry_str, local_entry)); + } else if file_type.is_file() { + let transfer_result: Result<(), String> = async { + let mut remote_file = sftp + .open(&remote_entry_str) + .await + .map_err(|e| format!("SFTP open error: {}", e))?; + let mut local_file = tokio::fs::File::create(&local_entry) + .await + .map_err(|e| format!("File create error: {}", e))?; + let mut buffer = vec![0u8; MAX_BUFF_SIZE]; + loop { + let n = remote_file + .read(&mut buffer) + .await + .map_err(|e| format!("File read error: {}", e))?; + if n == 0 { + break; + } + local_file + .write_all(&buffer[..n]) + .await + .map_err(|e| format!("File write error: {}", e))?; + } + local_file + .flush() + .await + .map_err(|e| format!("Flush error: {}", e))?; + #[cfg(unix)] + if preserve_permissions { + if let Some(perm) = sftp + .metadata(&remote_entry_str) + .await + .ok() + .and_then(|m| m.permissions) + { + use std::os::unix::fs::PermissionsExt; + let _ = tokio::fs::set_permissions( + &local_entry, + std::fs::Permissions::from_mode(perm & 0o7777), + ) + .await; + } + } + Ok(()) + } + .await; + match transfer_result { + Ok(_) => transferred.push(remote_entry_str), + Err(e) => { + if fail_fast { + return Err(PyRuntimeError::new_err(e)); + } + failed.push(remote_entry_str); + } + } + } + } + } + + Ok((transferred, failed)) + } + pub(crate) fn create( host: String, username: Option, @@ -755,7 +1172,58 @@ impl AsyncConnection { }) } - #[pyo3(signature = (pty=None))] + /// Uploads a local directory recursively to a remote path over SFTP. + /// Returns a tuple of (transferred_files, failed_files). + /// If `fail_fast` is true, the first transfer error raises an exception immediately. + #[pyo3(signature = (local_path, remote_path, follow_symlinks=true, preserve_permissions=true, fail_fast=false))] + fn sftp_put_dir<'p>( + &self, + py: Python<'p>, + local_path: String, + remote_path: String, + follow_symlinks: bool, + preserve_permissions: bool, + fail_fast: bool, + ) -> PyResult> { + let conn = self.clone(); + pyo3_async_runtimes::tokio::future_into_py(py, async move { + conn.sftp_put_dir_async( + local_path, + remote_path, + follow_symlinks, + preserve_permissions, + fail_fast, + ) + .await + }) + } + + /// Downloads a remote directory recursively to a local path over SFTP. + /// Returns a tuple of (transferred_files, failed_files). + /// If `fail_fast` is true, the first transfer error raises an exception immediately. + #[pyo3(signature = (remote_path, local_path, follow_symlinks=true, preserve_permissions=true, fail_fast=false))] + fn sftp_get_dir<'p>( + &self, + py: Python<'p>, + remote_path: String, + local_path: String, + follow_symlinks: bool, + preserve_permissions: bool, + fail_fast: bool, + ) -> PyResult> { + let conn = self.clone(); + pyo3_async_runtimes::tokio::future_into_py(py, async move { + conn.sftp_get_dir_async( + remote_path, + local_path, + follow_symlinks, + preserve_permissions, + fail_fast, + ) + .await + }) + } + fn shell<'p>(&self, py: Python<'p>, pty: Option) -> PyResult> { let session_arc = self.session.clone(); let pty = pty.unwrap_or(false); diff --git a/src/connection.rs b/src/connection.rs index e7f2bd5..e72cd12 100644 --- a/src/connection.rs +++ b/src/connection.rs @@ -67,7 +67,7 @@ use std::path::Path; use pyo3::exceptions::{PyIOError, PyTimeoutError}; -const MAX_BUFF_SIZE: usize = 65536; +pub(crate) const MAX_BUFF_SIZE: usize = 65536; create_exception!( connection, AuthenticationError, @@ -184,6 +184,26 @@ impl SSHResult { /// * `local_path`: The path to the file on the local system. /// * `remote_path`: The path to save the file on the remote system. /// +/// ### `sftp_put_dir` +/// +/// Uploads a local directory recursively to a remote path over SFTP. Returns (transferred_files, failed_files). It takes the following parameters: +/// +/// * `local_path`: The local directory to upload. +/// * `remote_path`: The remote destination path. +/// * `follow_symlinks`: Whether to follow symlinks (default: true). +/// * `preserve_permissions`: Whether to preserve file permissions (default: true). +/// * `fail_fast`: Whether to raise an exception on the first error instead of collecting failures (default: false). +/// +/// ### `sftp_get_dir` +/// +/// Downloads a remote directory recursively to a local path over SFTP. Returns (transferred_files, failed_files). It takes the following parameters: +/// +/// * `remote_path`: The remote directory to download. +/// * `local_path`: The local destination path. +/// * `follow_symlinks`: Whether to follow symlinks (default: true). +/// * `preserve_permissions`: Whether to preserve file permissions (default: true). +/// * `fail_fast`: Whether to raise an exception on the first error instead of collecting failures (default: false). +/// /// ### `shell` /// /// Creates an `InteractiveShell` instance. It takes the following parameter: @@ -561,6 +581,350 @@ impl Connection { Ok(()) } + /// Uploads a local directory recursively to a remote path over SFTP. + /// Returns a tuple of (transferred_files, failed_files), where each is a list of local file paths. + /// If `fail_fast` is true, the first transfer error raises an exception immediately. + #[pyo3(signature = (local_path, remote_path, follow_symlinks=true, preserve_permissions=true, fail_fast=false))] + fn sftp_put_dir( + &mut self, + local_path: String, + remote_path: String, + follow_symlinks: bool, + preserve_permissions: bool, + fail_fast: bool, + ) -> PyResult<(Vec, Vec)> { + let mut transferred: Vec = Vec::new(); + let mut failed: Vec = Vec::new(); + + // Ensure remote base directory (and all missing parents) exist — like mkdir -p + { + let parts: Vec<&str> = remote_path.split('/').filter(|s| !s.is_empty()).collect(); + let is_absolute = remote_path.starts_with('/'); + let mut current = if is_absolute { + String::from("/") + } else { + String::new() + }; + for part in &parts { + if !current.is_empty() && !current.ends_with('/') { + current.push('/'); + } + current.push_str(part); + if self.sftp().stat(Path::new(¤t)).is_err() { + if let Err(e) = self.sftp().mkdir(Path::new(¤t), 0o755) { + let msg = format!("Failed to create remote directory '{}': {}", current, e); + if fail_fast { + return Err(PyErr::new::(msg)); + } + failed.push(local_path.clone()); + return Ok((transferred, failed)); + } + } + } + } + + // Stack for depth-first traversal: (local_dir, remote_dir_str). + // Remote paths are kept as POSIX strings to avoid OS-native path separators on Windows. + let mut dirs_to_process: Vec<(std::path::PathBuf, String)> = + vec![(std::path::PathBuf::from(&local_path), remote_path.clone())]; + + while let Some((local_dir, remote_dir)) = dirs_to_process.pop() { + let entries = match std::fs::read_dir(&local_dir) { + Ok(e) => e, + Err(e) => { + let msg = format!( + "Failed to read local directory '{}': {}", + local_dir.display(), + e + ); + if fail_fast { + return Err(PyErr::new::(msg)); + } + failed.push(local_dir.to_string_lossy().to_string()); + continue; + } + }; + + for entry in entries { + let entry = match entry { + Ok(e) => e, + Err(e) => { + if fail_fast { + return Err(PyErr::new::(format!( + "Directory entry error: {}", + e + ))); + } + continue; + } + }; + + let local_entry = entry.path(); + let local_entry_str = local_entry.to_string_lossy().to_string(); + let file_name_str = entry.file_name().to_string_lossy().to_string(); + let remote_entry_str = format!("{}/{}", remote_dir, file_name_str); + + let metadata = match if follow_symlinks { + std::fs::metadata(&local_entry) + } else { + std::fs::symlink_metadata(&local_entry) + } { + Ok(m) => m, + Err(e) => { + if fail_fast { + return Err(PyErr::new::(format!( + "Metadata error for '{}': {}", + local_entry_str, e + ))); + } + failed.push(local_entry_str); + continue; + } + }; + + if metadata.is_dir() { + #[cfg(unix)] + let mode = if preserve_permissions { + use std::os::unix::fs::PermissionsExt; + metadata.permissions().mode() as i32 + } else { + 0o755i32 + }; + #[cfg(not(unix))] + let mode = 0o755i32; + if self.sftp().stat(Path::new(&remote_entry_str)).is_err() { + match self.sftp().mkdir(Path::new(&remote_entry_str), mode) { + Ok(_) => {} + Err(e) => { + if fail_fast { + return Err(PyErr::new::(format!( + "Failed to create remote directory '{}': {}", + remote_entry_str, e + ))); + } + failed.push(local_entry_str); + continue; + } + } + } + dirs_to_process.push((local_entry, remote_entry_str)); + } else if metadata.is_file() { + let result: Result<(), String> = (|| { + let mut local_file = std::fs::File::open(&local_entry) + .map_err(|e| format!("File open error: {}", e))?; + let file_size = metadata.len(); + let mut remote_file = self + .sftp() + .create(Path::new(&remote_entry_str)) + .map_err(|e| format!("Remote file creation error: {}", e))?; + let buf_size = (file_size as usize).min(MAX_BUFF_SIZE); + // Fall back to MAX_BUFF_SIZE for empty files (size 0) so the read loop can still run. + let buf_size = if buf_size == 0 { + MAX_BUFF_SIZE + } else { + buf_size + }; + let mut buffer = vec![0u8; buf_size]; + loop { + let n = local_file + .read(&mut buffer) + .map_err(|e| format!("File read error: {}", e))?; + if n == 0 { + break; + } + remote_file + .write_all(&buffer[..n]) + .map_err(|e| format!("Remote write error: {}", e))?; + } + remote_file + .close() + .map_err(|e| format!("Close error: {}", e))?; + #[cfg(unix)] + if preserve_permissions { + use std::os::unix::fs::PermissionsExt; + let mode = metadata.permissions().mode(); + let _ = self.sftp().setstat( + Path::new(&remote_entry_str), + ssh2::FileStat { + perm: Some(mode), + size: None, + uid: None, + gid: None, + atime: None, + mtime: None, + }, + ); + } + Ok(()) + })(); + match result { + Ok(_) => transferred.push(local_entry_str), + Err(e) => { + if fail_fast { + return Err(PyErr::new::(e)); + } + failed.push(local_entry_str); + } + } + } + // Symlinks with follow_symlinks=false are skipped + } + } + + Ok((transferred, failed)) + } + + /// Downloads a remote directory recursively to a local path over SFTP. + /// Returns a tuple of (transferred_files, failed_files), where each is a list of remote file paths. + /// If `fail_fast` is true, the first transfer error raises an exception immediately. + #[pyo3(signature = (remote_path, local_path, follow_symlinks=true, preserve_permissions=true, fail_fast=false))] + fn sftp_get_dir( + &mut self, + remote_path: String, + local_path: String, + follow_symlinks: bool, + preserve_permissions: bool, + fail_fast: bool, + ) -> PyResult<(Vec, Vec)> { + let mut transferred: Vec = Vec::new(); + let mut failed: Vec = Vec::new(); + + // Ensure local base directory exists + if let Err(e) = std::fs::create_dir_all(&local_path) { + let msg = format!("Failed to create local directory '{}': {}", local_path, e); + if fail_fast { + return Err(PyErr::new::(msg)); + } + failed.push(remote_path); + return Ok((transferred, failed)); + } + + // Stack for depth-first traversal: (remote_dir_str, local_dir). + // Remote paths are kept as POSIX strings to avoid OS-native path separators on Windows. + let mut dirs_to_process: Vec<(String, std::path::PathBuf)> = + vec![(remote_path.clone(), std::path::PathBuf::from(&local_path))]; + + while let Some((remote_dir, local_dir)) = dirs_to_process.pop() { + let remote_dir_str = remote_dir.clone(); + let entries = match self.sftp().readdir(Path::new(&remote_dir_str)) { + Ok(e) => e, + Err(e) => { + let msg = format!( + "Failed to read remote directory '{}': {}", + remote_dir_str, e + ); + if fail_fast { + return Err(PyErr::new::(msg)); + } + failed.push(remote_dir_str); + continue; + } + }; + + for (entry_name, stat) in entries { + let file_name = match entry_name.file_name() { + Some(n) => n.to_os_string(), + None => { + if fail_fast { + return Err(PyErr::new::("Invalid entry path")); + } + continue; + } + }; + let file_name_str = file_name.to_string_lossy(); + let local_entry = local_dir.join(&file_name); + // Build remote path with POSIX separator to stay cross-platform. + let remote_entry_str = format!("{}/{}", remote_dir, file_name_str); + + // When follow_symlinks=true, resolve symlinks on remote via stat() + let resolved_stat = if follow_symlinks && stat.file_type().is_symlink() { + self.sftp() + .stat(Path::new(&remote_entry_str)) + .unwrap_or(stat) + } else { + stat + }; + + if resolved_stat.file_type().is_symlink() { + // follow_symlinks=false: skip symlinks + continue; + } else if resolved_stat.is_dir() { + match std::fs::create_dir_all(&local_entry) { + Ok(_) => {} + Err(e) => { + if fail_fast { + return Err(PyErr::new::(format!( + "Failed to create local directory '{}': {}", + local_entry.display(), + e + ))); + } + failed.push(remote_entry_str); + continue; + } + } + #[cfg(unix)] + if preserve_permissions { + if let Some(perm) = resolved_stat.perm { + use std::os::unix::fs::PermissionsExt; + let _ = std::fs::set_permissions( + &local_entry, + std::fs::Permissions::from_mode(perm & 0o7777), + ); + } + } + dirs_to_process.push((remote_entry_str, local_entry)); + } else if resolved_stat.is_file() { + let result: Result<(), String> = (|| { + let mut remote_file = BufReader::new( + self.sftp() + .open(Path::new(&remote_entry_str)) + .map_err(|e| format!("SFTP open error: {}", e))?, + ); + let local_file = std::fs::File::create(&local_entry) + .map_err(|e| format!("File create error: {}", e))?; + let mut writer = BufWriter::new(local_file); + let mut buffer = vec![0u8; MAX_BUFF_SIZE]; + loop { + let n = remote_file + .read(&mut buffer) + .map_err(|e| format!("File read error: {}", e))?; + if n == 0 { + break; + } + writer + .write_all(&buffer[..n]) + .map_err(|e| format!("File write error: {}", e))?; + } + writer.flush().map_err(|e| format!("Flush error: {}", e))?; + #[cfg(unix)] + if preserve_permissions { + if let Some(perm) = resolved_stat.perm { + use std::os::unix::fs::PermissionsExt; + let _ = std::fs::set_permissions( + &local_entry, + std::fs::Permissions::from_mode(perm & 0o7777), + ); + } + } + Ok(()) + })(); + match result { + Ok(_) => transferred.push(remote_entry_str), + Err(e) => { + if fail_fast { + return Err(PyErr::new::(e)); + } + failed.push(remote_entry_str); + } + } + } + } + } + + Ok((transferred, failed)) + } + // Copy a file from this connection to another connection #[pyo3(signature = (source_path, dest_conn, dest_path=None))] fn remote_copy( diff --git a/tests/test_async_connection.py b/tests/test_async_connection.py index 957b85c..658e707 100644 --- a/tests/test_async_connection.py +++ b/tests/test_async_connection.py @@ -66,6 +66,76 @@ async def test_async_sftp_write_data(run_test_server): assert "hello.txt" in files +@pytest.mark.asyncio +async def test_async_sftp_put_dir(run_test_server, tmp_path): + """Test that we can recursively upload a directory tree over async SFTP.""" + async with AsyncConnection("localhost", username="root", password="toor", port=8022) as conn: + # Create a local directory tree + src = tmp_path / "async_src_dir" + src.mkdir() + (src / "file1.txt").write_text("file1 content") + (src / "file2.txt").write_text("file2 content") + sub = src / "subdir" + sub.mkdir() + (sub / "nested.txt").write_text("nested content") + + # Upload to remote + transferred, failed = await conn.sftp_put_dir(str(src), "/root/async_test_put_dir") + + # Verify files exist on remote + remote_ls = (await conn.execute("find /root/async_test_put_dir -type f | sort")).stdout + assert "file1.txt" in remote_ls + assert "file2.txt" in remote_ls + assert "nested.txt" in remote_ls + expected_file_count = 3 + assert len(transferred) == expected_file_count + assert len(failed) == 0 + assert any("file1.txt" in p for p in transferred) + + # Verify nested file contents + content = await conn.sftp_read("/root/async_test_put_dir/subdir/nested.txt") + assert content == "nested content" + + # Cleanup + await conn.execute("rm -rf /root/async_test_put_dir") + + +@pytest.mark.asyncio +async def test_async_sftp_get_dir(run_test_server, tmp_path): + """Test that we can recursively download a directory tree over async SFTP.""" + async with AsyncConnection("localhost", username="root", password="toor", port=8022) as conn: + # Set up a remote directory tree + await conn.execute("mkdir -p /root/async_test_get_dir/subdir") + await conn.sftp_write_data("remote file 1", "/root/async_test_get_dir/file1.txt") + await conn.sftp_write_data("remote file 2", "/root/async_test_get_dir/file2.txt") + await conn.sftp_write_data("nested remote", "/root/async_test_get_dir/subdir/nested.txt") + + # Download to local + dest = tmp_path / "async_dest_dir" + transferred, failed = await conn.sftp_get_dir("/root/async_test_get_dir", str(dest)) + + # Verify local files + assert (dest / "file1.txt").read_text() == "remote file 1" + assert (dest / "file2.txt").read_text() == "remote file 2" + assert (dest / "subdir" / "nested.txt").read_text() == "nested remote" + expected_file_count = 3 + assert len(transferred) == expected_file_count + assert len(failed) == 0 + assert any("file1.txt" in p for p in transferred) + + # Cleanup + await conn.execute("rm -rf /root/async_test_get_dir") + + +@pytest.mark.asyncio +async def test_async_sftp_get_dir_fail_fast(run_test_server, tmp_path): + """Test that sftp_get_dir with fail_fast=True raises on error.""" + async with AsyncConnection("localhost", username="root", password="toor", port=8022) as conn: + dest = tmp_path / "async_dest_fail" + with pytest.raises(RuntimeError): + await conn.sftp_get_dir("/path/does/not/exist", str(dest), fail_fast=True) + + @pytest.mark.asyncio async def test_async_shell(run_test_server): async with ( diff --git a/tests/test_connection.py b/tests/test_connection.py index 0782603..9cc503a 100644 --- a/tests/test_connection.py +++ b/tests/test_connection.py @@ -212,6 +212,88 @@ def test_sftp_write_data(conn): assert read_text == "hello" +def test_sftp_put_dir(conn, tmp_path): + """Test that we can recursively upload a directory tree over SFTP.""" + # Create a local directory tree + src = tmp_path / "src_dir" + src.mkdir() + (src / "file1.txt").write_text("file1 content") + (src / "file2.txt").write_text("file2 content") + sub = src / "subdir" + sub.mkdir() + (sub / "nested.txt").write_text("nested content") + + # Upload to remote + transferred, failed = conn.sftp_put_dir(str(src), "/root/test_put_dir") + + # Verify files exist on remote + remote_ls = conn.execute("find /root/test_put_dir -type f | sort").stdout + assert "file1.txt" in remote_ls + assert "file2.txt" in remote_ls + assert "nested.txt" in remote_ls + expected_file_count = 3 + assert len(transferred) == expected_file_count + assert len(failed) == 0 + # Transferred list contains local paths + assert any("file1.txt" in p for p in transferred) + + # Verify nested file contents + content = conn.sftp_read("/root/test_put_dir/subdir/nested.txt") + assert content == "nested content" + + # Cleanup + conn.execute("rm -rf /root/test_put_dir") + + +def test_sftp_get_dir(conn, tmp_path): + """Test that we can recursively download a directory tree over SFTP.""" + # Set up a remote directory tree + conn.execute("mkdir -p /root/test_get_dir/subdir") + conn.sftp_write_data("remote file 1", "/root/test_get_dir/file1.txt") + conn.sftp_write_data("remote file 2", "/root/test_get_dir/file2.txt") + conn.sftp_write_data("nested remote", "/root/test_get_dir/subdir/nested.txt") + + # Download to local + dest = tmp_path / "dest_dir" + transferred, failed = conn.sftp_get_dir("/root/test_get_dir", str(dest)) + + # Verify local files + assert (dest / "file1.txt").read_text() == "remote file 1" + assert (dest / "file2.txt").read_text() == "remote file 2" + assert (dest / "subdir" / "nested.txt").read_text() == "nested remote" + expected_file_count = 3 + assert len(transferred) == expected_file_count + assert len(failed) == 0 + # Transferred list contains remote paths + assert any("file1.txt" in p for p in transferred) + + # Cleanup + conn.execute("rm -rf /root/test_get_dir") + + +def test_sftp_put_dir_fail_fast(conn, tmp_path): + """Test that sftp_put_dir with fail_fast=True raises on error.""" + src = tmp_path / "src_fail" + src.mkdir() + (src / "ok.txt").write_text("ok") + + # Pre-create a *file* at the target path so that attempting to create a + # sub-directory inside it fails, even after mkdir-p creates its parents. + conn.sftp_write_data("blocking", "/root/put_dir_fail_target") + try: + with pytest.raises(OSError, match=r"(?i)(failed|error|no such)"): + conn.sftp_put_dir(str(src), "/root/put_dir_fail_target/deep", fail_fast=True) + finally: + conn.execute("rm -f /root/put_dir_fail_target") + + +def test_sftp_get_dir_fail_fast(conn, tmp_path): + """Test that sftp_get_dir with fail_fast=True raises on error.""" + dest = tmp_path / "dest_fail" + with pytest.raises(OSError, match=r"(?i)(failed|error|no such)"): + conn.sftp_get_dir("/path/does/not/exist", str(dest), fail_fast=True) + + @pytest.mark.skip("non-text files are not supported by sftp") def test_non_utf8_sftp(conn): """Test that we can copy a non-text file to the server and read it back."""