diff --git a/src/args.rs b/src/args.rs index 7c46f27..8084f5d 100644 --- a/src/args.rs +++ b/src/args.rs @@ -1,12 +1,5 @@ use std::num::NonZeroU32; -use clap::{Parser, ArgGroup}; - -#[derive(Clone, Debug)] -pub enum CLIAction { - DownloadUrl(String), - ResolveZippyUrl(String), - UrlList(String), -} +use clap::Parser; #[derive(Parser, Clone, Debug)] #[clap( @@ -14,28 +7,22 @@ pub enum CLIAction { about, long_about = None, name = "FFDL - Fast File Downloader", - group( - ArgGroup::new("action") - .required(true) - .args(&["listfile", "download", "zippy-resolve"]) - ) )] pub struct CLIArgs { #[clap( - short, - long, + short = 'o', + long = "outdir", value_name = "OUTPUT DIR", default_value = "./", - help = "Set the output directory. The directory will be created if it doesn't exit yet" + help = "Set the output directory. The directory will be created if it doesn't exit yet", )] pub outdir: String, #[clap( - short, + short = 'i', long = "into-file", value_name = "FILENAME", help = "Force filename. This only works for single file downloads", - requires = "download" )] pub into_file: Option, @@ -45,50 +32,42 @@ pub struct CLIArgs { value_name = "NUMBER OF CONCURRENT FILE DOWNLOADS", default_value = "1", help = "Specify the number of concurrent downloads", - requires = "listfile" )] pub file_count: NonZeroU32, #[clap( - short, + short = 'c', long = "connections", value_name = "NUMBER OF CONCURRENT CONNECTIONS", default_value = "1", help = "The number concurrent connections per file download. \ Downloads might fail when the number of connections is too high. \ Files started with multiple connections currently can't be continued. \ - NOTE: This will likely cause IO bottlenecks on HDDs" - )] - pub conn_count: NonZeroU32, + NOTE: This will likely cause IO bottlenecks on HDDs", + )] + pub conn_count: NonZeroU32, #[clap( - short, - long, + short = 'z', + long = "zippy", help = "The provided URLs are zippyshare URLs and need to be \ - resolved to direct download urls" + resolved to direct download urls", )] pub zippy: bool, - #[clap( + #[clap( short = 'l', long = "listfile", value_name = "URL LISTFILE", help = "Download all files from the specified url list file", )] - pub listfile: Option, + pub listfile: Vec, #[clap( short = 'd', long = "download", value_name = "URL", - help = "Download only the one file fromn the specified url", + help = "Download only the one file from the specified url", )] - pub download: Option, - - #[clap( - long = "zippy-resolve", - value_name = "ZIPPYSHARE URL", - help = "Resolve the zippy share url to the direct download url", - )] - pub zippy_resolve: Option, -} \ No newline at end of file + pub download: Vec, +} diff --git a/src/main.rs b/src/main.rs index c8e042d..b77a858 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,14 +1,17 @@ use std::{ - collections::VecDeque, io::BufRead, path::Path, process::exit, sync::Arc, sync::Mutex, + collections::VecDeque, + path::{Path, PathBuf}, + process::exit, + sync::Arc, time::SystemTime, }; use clap::Parser; use futures::future::join_all; -use tokio::sync::mpsc; +use tokio::sync::{mpsc, Mutex}; use crate::{ - args::{CLIAction, CLIArgs}, + args::CLIArgs, dlreport::{DlReport, DlReporter, DlStatus}, errors::ResBE, }; @@ -23,193 +26,70 @@ mod zippy; async fn main() -> ResBE<()> { let args = CLIArgs::parse(); - let action = match (&args.listfile, &args.download, &args.zippy_resolve) { - (Some(listfile), None, None) => CLIAction::UrlList(listfile.to_string()), - (None, Some(url), None) => CLIAction::DownloadUrl(url.to_string()), - (None, None, Some(zippy_url)) => CLIAction::ResolveZippyUrl(zippy_url.to_string()), - _ => unreachable!(), - }; - - let urls = match action { - CLIAction::DownloadUrl(url) => vec![url.clone()], - CLIAction::UrlList(listfile) => read_urls_from_listfile(&listfile).await, - CLIAction::ResolveZippyUrl(url) => resolve_zippy_url(&url).await, - }; + let mut urls = args.download.clone(); + for file in args.listfile.iter() { + match read_urls_from_listfile(file).await { + Ok(listfile_urls) => urls.extend(listfile_urls), + Err(_) => { + eprintln!("Failed to read urls from file: {}", file); + exit(1); + } + } + } download_multiple(args, urls).await } -async fn resolve_zippy_url(url: &str) -> ! { - let resolved_url = zippy::resolve_link(&url).await.unwrap_or_else(|_| { - println!("Zippyshare link could not be resolved"); - exit(1); - }); - - println!("{}", resolved_url); - exit(0); -} - -async fn read_urls_from_listfile(listfile: &str) -> Vec { - let p_listfile = Path::new(&listfile); - - if !p_listfile.is_file() { - eprintln!("Listfile '{}' does not exist!", &listfile); - exit(1); - } - - let ifile = std::fs::File::open(p_listfile).unwrap(); - - std::io::BufReader::new(ifile) +async fn read_urls_from_listfile(listfile: &str) -> ResBE> { + let text = tokio::fs::read_to_string(listfile).await?; + let urls = text .lines() - .map(|l| l.unwrap()) - .filter(|url| url.len() > 0 && !url.starts_with("#")) - .collect() + .map(str::trim) + .filter(|line| !line.is_empty() && !line.starts_with('#')) + .map(str::to_string) + .collect(); + Ok(urls) } async fn download_multiple(cli_args: CLIArgs, raw_urls: Vec) -> ResBE<()> { - let outdir = cli_args.outdir; - let outdir = Path::new(&outdir); + let outdir = Path::new(&cli_args.outdir); + let num_urls = raw_urls.len(); let parallel_file_count = cli_args.file_count.get(); let conn_count = cli_args.conn_count.get(); let zippy = cli_args.zippy; - let urls = Arc::new(Mutex::new(VecDeque::<(u32, String)>::new())); + let urls = Arc::new(Mutex::new(VecDeque::<(usize, String)>::new())); - raw_urls.iter().enumerate().for_each(|(i, url)| { - urls.lock().unwrap().push_back((i as u32, url.clone())); - }); + urls.lock().await.extend(raw_urls.into_iter().enumerate()); if !outdir.exists() { - if let Err(_e) = std::fs::create_dir_all(outdir) { - eprintln!( - "Error creating output directory '{}'", - outdir.to_str().unwrap() - ); + if let Err(_e) = tokio::fs::create_dir_all(outdir).await { + eprintln!("Error creating output directory '{}'", outdir.display()); exit(1); } } - let t_start = SystemTime::now(); - - let mut joiners = Vec::new(); - let (tx, rx) = mpsc::unbounded_channel::(); - for _offset in 0..parallel_file_count { - let tx = tx.clone(); - let outdir = outdir.to_owned(); - let arg_filename = cli_args.into_file.clone(); + let t_start = SystemTime::now(); - let urls = urls.clone(); - - joiners.push(tokio::task::spawn(async move { - loop { - let (global_url_index, url) = match urls.lock().unwrap().pop_front() { - Some(it) => it, - None => break, - }; - - let tx = tx.clone(); - - let rep = DlReporter::new(global_url_index, tx); - - let url = if zippy { - match zippy::resolve_link(&url).await { - Ok(url) => url, - Err(_e) => { - rep.send(DlStatus::Message(format!( - "Zippyshare link could not be resolved: {}", - url - ))); - continue; - } - } - } else { - url.to_string() - }; - - let file_name = arg_filename - .clone() - .unwrap_or_else(|| download::url_to_filename(&url)); - - let into_file = outdir - .join(Path::new(&file_name)) - .to_str() - .unwrap() - .to_string(); - let path_into_file = Path::new(&into_file); - - let (filesize, range_supported) = - match download::http_get_filesize_and_range_support(&url).await { - Ok((filesize, range_supported)) => (filesize, range_supported), - Err(_e) => { - rep.send(DlStatus::Message(format!( - "Error while querying metadata: {}", - url - ))); - continue; - } - }; - - // If file with same name is present locally, check filesize - if path_into_file.exists() { - let local_filesize = std::fs::metadata(path_into_file).unwrap().len(); - - if filesize == local_filesize { - rep.send(DlStatus::Message(format!( - "Skipping file '{}': already present", - &file_name - ))); - rep.send(DlStatus::Skipped); - continue; - } else { - rep.send(DlStatus::Message(format!( - "Replacing file '{}': present but not completed", - &file_name - ))); - } - } - - if conn_count == 1 { - if let Err(_e) = - download::download_feedback(&url, &into_file, rep.clone(), Some(filesize)) - .await - { - rep.send(DlStatus::DoneErr { - filename: file_name.to_string(), - }); - } - } else { - if !range_supported { - rep.send(DlStatus::Message(format!( - "Error Server does not support range header: {}", - url - ))); - continue; - } - - if let Err(_e) = download::download_feedback_multi( - &url, - &into_file, - rep.clone(), - conn_count, - Some(filesize), - ) - .await - { - rep.send(DlStatus::DoneErr { - filename: file_name.to_string(), - }); - } - }; - } - })) - } + let joiners = (0..parallel_file_count) + .map(|_| { + tokio::task::spawn(download_job( + urls.clone(), + tx.clone(), + conn_count, + zippy, + outdir.to_owned(), + cli_args.into_file.clone(), + )) + }) + .collect::>(); drop(tx); - dlreport::watch_and_print_reports(rx, raw_urls.len() as i32).await?; + dlreport::watch_and_print_reports(rx, num_urls as i32).await?; join_all(joiners).await; @@ -217,3 +97,112 @@ async fn download_multiple(cli_args: CLIArgs, raw_urls: Vec) -> ResBE<() Ok(()) } + +async fn download_job( + urls: Arc>>, + tx: mpsc::UnboundedSender, + conn_count: u32, + zippy: bool, + outdir: PathBuf, + arg_filename: Option, +) { + loop { + let (global_url_index, url) = match urls.lock().await.pop_front() { + Some(it) => it, + None => break, + }; + + let tx = tx.clone(); + + let rep = DlReporter::new(global_url_index as u32, tx); + + let url = if zippy { + match zippy::resolve_link(&url).await { + Ok(url) => url, + Err(_e) => { + rep.send(DlStatus::Message(format!( + "Zippyshare link could not be resolved: {}", + url + ))); + continue; + } + } + } else { + url.to_string() + }; + + let file_name = arg_filename + .clone() + .unwrap_or_else(|| download::url_to_filename(&url)); + + let into_file = outdir + .join(Path::new(&file_name)) + .to_str() + .unwrap() + .to_string(); + let path_into_file = Path::new(&into_file); + + let (filesize, range_supported) = + match download::http_get_filesize_and_range_support(&url).await { + Ok((filesize, range_supported)) => (filesize, range_supported), + Err(_e) => { + rep.send(DlStatus::Message(format!( + "Error while querying metadata: {}", + url + ))); + continue; + } + }; + + // If file with same name is present locally, check filesize + if path_into_file.exists() { + let local_filesize = std::fs::metadata(path_into_file).unwrap().len(); + + if filesize == local_filesize { + rep.send(DlStatus::Message(format!( + "Skipping file '{}': already present", + &file_name + ))); + rep.send(DlStatus::Skipped); + continue; + } else { + rep.send(DlStatus::Message(format!( + "Replacing file '{}': present but not completed", + &file_name + ))); + } + } + + if conn_count == 1 { + if let Err(_e) = + download::download_feedback(&url, &into_file, rep.clone(), Some(filesize)).await + { + rep.send(DlStatus::DoneErr { + filename: file_name.to_string(), + }); + } + } else { + if !range_supported { + rep.send(DlStatus::Message(format!( + "Error Server does not support range header: {}", + url + ))); + continue; + } + + if let Err(_e) = download::download_feedback_multi( + &url, + &into_file, + rep.clone(), + conn_count, + Some(filesize), + ) + .await + { + rep.send(DlStatus::DoneErr { + filename: file_name.to_string(), + }); + } + }; + } +}