sproto 0.1.0

Rust client for the Synology Drive sync protocol
Documentation
#[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 to download
    remote_path: String,

    /// Local destination path
    local_path: PathBuf,

    /// Parallel downloads
    #[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)
}