#[path = "support/mod.rs"]
mod support;
use anyhow::{Context, Result};
use clap::Parser;
use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
use std::path::PathBuf;
use std::sync::Arc;
use tokio::sync::Semaphore;
use sproto::actions::list::{FileType, NodeInfo};
#[derive(Parser)]
#[command(
name = "sproto-dl-dir",
about = "Download a directory from Synology Drive"
)]
struct Args {
remote_path: String,
local_path: PathBuf,
#[arg(long, default_value = "5", value_parser = clap::value_parser!(u8).range(1..))]
concurrency: u8,
#[command(flatten)]
conn: support::ConnectionArgs,
}
#[tokio::main]
async fn main() -> Result<()> {
support::init_tracing();
let args = Args::parse();
let client = args.conn.connect().await?;
let view_id = args.conn.resolve_view(&client).await?;
eprintln!("Enumerating {}...", args.remote_path);
let files = enumerate_recursive(&client, view_id, &args.remote_path).await?;
eprintln!("Found {} files to download", files.len());
if files.is_empty() {
return Ok(());
}
let concurrency = usize::from(args.concurrency);
client.warm_pool(concurrency).await?;
let sem = Arc::new(Semaphore::new(concurrency));
let mp = MultiProgress::new();
let total_bytes: u64 = files.iter().map(|(n, _)| n.file_size).sum();
let overall = mp.add(ProgressBar::new(total_bytes));
overall.set_style(
#[allow(
clippy::literal_string_with_formatting_args,
reason = "indicatif template syntax"
)]
ProgressStyle::default_bar()
.template("{msg} [{bar:40}] {bytes}/{total_bytes} ({bytes_per_sec})")
.unwrap()
.progress_chars("=> "),
);
overall.set_message("Total");
let root = args.remote_path.trim_end_matches('/');
let mut handles = vec![];
for (node, remote_path) in &files {
let permit = sem.clone().acquire_owned().await?;
let client = client.clone();
let relative = remote_path
.strip_prefix(root)
.unwrap_or(remote_path)
.trim_start_matches('/');
let local = args.local_path.join(relative);
let file_id = node.file_id.clone();
let file_size = node.file_size;
let name = node.name.clone();
let overall = overall.clone();
let mp = mp.clone();
handles.push(tokio::spawn(async move {
let _permit = permit;
let pb = mp.insert_before(&overall, ProgressBar::new(file_size));
pb.set_style(
#[allow(
clippy::literal_string_with_formatting_args,
reason = "indicatif template syntax"
)]
ProgressStyle::default_bar()
.template(" {msg:>20} [{bar:30}] {bytes}/{total_bytes}")
.unwrap()
.progress_chars("=> "),
);
pb.set_message(name);
let mut writer = support::ProgressWriter::new(local, pb.clone()).with_overall(overall);
let result = client.download_to(&file_id, &mut writer).await;
pb.finish_and_clear();
result
}));
}
let mut errors = 0;
for h in handles {
if let Err(e) = h.await? {
eprintln!("Download error: {e}");
errors += 1;
}
}
overall.finish_with_message("Done");
if errors > 0 {
anyhow::bail!("{errors} file(s) failed to download");
}
Ok(())
}
async fn enumerate_recursive(
client: &sproto::Client,
view_id: u64,
path: &str,
) -> Result<Vec<(NodeInfo, String)>> {
let mut result = Vec::new();
let mut dirs_to_visit = vec![path.to_string()];
while let Some(dir_path) = dirs_to_visit.pop() {
let nodes = client
.list_dir(view_id, &dir_path)
.await
.with_context(|| format!("failed to list '{dir_path}'"))?;
for node in nodes {
if node.is_removed {
continue;
}
let parent = dir_path.trim_end_matches('/');
let full_path = format!("{parent}/{}", node.name);
match node.file_type {
FileType::Dir => {
dirs_to_visit.push(full_path);
}
FileType::File => {
result.push((node, full_path));
}
FileType::Symlink => {}
}
}
}
Ok(result)
}