From 2d2aaef9f822e874a6f49e10e519043675121929 Mon Sep 17 00:00:00 2001 From: Daniel M Date: Fri, 2 Apr 2021 15:58:16 +0200 Subject: [PATCH] Add into-file cli argument - Added option to set the filename for single file downloads - Changed the number cli arguments to unsigned, since they can only be positive anyways - Combined the cli arguments into a struct --- src/dlreport.rs | 10 +-- src/download.rs | 2 +- src/main.rs | 171 +++++++++++++++++++++++++++++++++++------------- 3 files changed, 132 insertions(+), 51 deletions(-) diff --git a/src/dlreport.rs b/src/dlreport.rs index b69cf95..e77f570 100644 --- a/src/dlreport.rs +++ b/src/dlreport.rs @@ -33,18 +33,18 @@ pub enum DlStatus { #[derive(Clone, Debug)] pub struct DlReport { - pub id: i32, + pub id: u32, pub status: DlStatus } #[derive(Clone)] pub struct DlReporter { - id: i32, + id: u32, transmitter: mpsc::UnboundedSender } impl DlReporter { - pub fn new(id: i32, transmitter: mpsc::UnboundedSender) -> DlReporter { + pub fn new(id: u32, transmitter: mpsc::UnboundedSender) -> DlReporter { DlReporter { id: id, transmitter: transmitter @@ -83,7 +83,7 @@ impl InfoHolder { } -fn print_accumulated_report(statuses: & HashMap, msg_queue: &mut VecDeque, moved_lines: u16) -> ResBE { +fn print_accumulated_report(statuses: & HashMap, msg_queue: &mut VecDeque, moved_lines: u16) -> ResBE { let mut dl_speed_sum = 0.0; execute!( @@ -142,7 +142,7 @@ fn print_accumulated_report(statuses: & HashMap, msg_queue: &mu pub async fn watch_and_print_reports(mut receiver: mpsc::UnboundedReceiver) -> ResBE<()> { - let mut statuses: HashMap = HashMap::new(); + let mut statuses: HashMap = HashMap::new(); let mut moved_lines = 0; let mut msg_queue = VecDeque::new(); diff --git a/src/download.rs b/src/download.rs index 3e6ecdb..eab3438 100644 --- a/src/download.rs +++ b/src/download.rs @@ -229,7 +229,7 @@ pub async fn download_feedback_chunks(url: &str, into_file: &str, rep: DlReporte // This will spin up multiple tasks that and manage the status updates for them. // The combined status will be reported back to the caller -pub async fn download_feedback_multi(url: &str, into_file: &str, rep: DlReporter, conn_count: i32, content_length: Option) -> ResBE<()> { +pub async fn download_feedback_multi(url: &str, into_file: &str, rep: DlReporter, conn_count: u32, content_length: Option) -> ResBE<()> { let content_length = match content_length { Some(it) => it, diff --git a/src/main.rs b/src/main.rs index a8ee53e..d75af94 100644 --- a/src/main.rs +++ b/src/main.rs @@ -15,6 +15,26 @@ mod errors; mod dlreport; +#[derive(Clone, Debug)] +enum CLIAction { + DownloadUrl(String), + ResolveZippyUrl(String), + UrlList(String), + None +} + +#[derive(Clone, Debug)] +struct CLIArguments { + outdir: String, + into_file: Option, + file_count: u32, + conn_count: u32, + zippy: bool, + action: CLIAction, + urls: Vec +} + + #[tokio::main] async fn main() -> ResBE<()> { @@ -30,6 +50,15 @@ async fn main() -> ResBE<()> { .help("Set the output directory. The directory will be created \ if it doesn't exit yet") ) + .arg( + Arg::with_name("into_file") + .short("i") + .long("into-file") + .value_name("FILENAME") + .takes_value(true) + .requires("download") + .help("Force filename. This only works for single file downloads") + ) .arg( Arg::with_name("file_count") .short("n") @@ -45,9 +74,10 @@ async fn main() -> ResBE<()> { .value_name("NUMBER OF CONNECTIONS") .takes_value(true) .help("The number concurrent connections per file download. \ - Downloads might fail when the number of connections is too high. \ - Files started with multiple connections can't be continued. \ - NOTE: This will likely cause IO bottlenecks on HDDs") + Downloads might fail when the number of connections is \ + too high. Files started with multiple connections can't \ + be continued. NOTE: This will likely cause IO \ + bottlenecks on HDDs") ) .arg( Arg::with_name("zippyshare") @@ -94,12 +124,17 @@ async fn main() -> ResBE<()> { None => "./" }; + let into_file = match arguments.value_of("into_file") { + Some(it) => Some(it.to_string()), + None => None + }; + let file_count = match arguments.value_of("file_count") { Some(it) => it, None => "1" }; - let file_count: i32 = match file_count.parse() { + let file_count: u32 = match file_count.parse() { Ok(it) => it, Err(_) => { eprintln!("Invalid value for num-files: {}", file_count); @@ -117,7 +152,7 @@ async fn main() -> ResBE<()> { None => "1" }; - let conn_count: i32 = match conn_count.parse() { + let conn_count: u32 = match conn_count.parse() { Ok(it) => it, Err(_) => { eprintln!("Invalid value for connections: {}", conn_count); @@ -133,56 +168,94 @@ async fn main() -> ResBE<()> { let is_zippy = arguments.is_present("zippyshare"); + let action = + if let Some(listfile) = arguments.value_of("listfile") { + CLIAction::UrlList ( + listfile.to_string() + ) + } else if let Some(download_url) = arguments.value_of("download") { + CLIAction::DownloadUrl( + download_url.to_string() + ) + } else if let Some(resolve_url) = arguments.value_of("zippy-resolve") { + CLIAction::ResolveZippyUrl( + resolve_url.to_string() + ) + } + else { + CLIAction::None + }; + + + let mut cli_args = CLIArguments { + outdir: outdir.to_string(), + into_file: into_file, + file_count: file_count, + conn_count: conn_count, + zippy: is_zippy, + action: action, + urls: Vec::new() + }; + // Evaluate and execute the requested action. The 3 different actions are // mutally exclusive, so only one of them will be executed - if let Some(s_listfile) = arguments.value_of("listfile") { + match &cli_args.action { - let listfile = Path::new(s_listfile); + CLIAction::UrlList(listfile) => { - if !listfile.is_file() { - eprintln!("Listfile '{}' does not exist!", s_listfile); - exit(1); - } + let p_listfile = Path::new(listfile); - let ifile = std::fs::File::open(listfile)?; - - let urls: Vec = std::io::BufReader::new(ifile) - .lines() - .map(|l| l.unwrap()) - .filter(|url| url.len() > 0 && !url.starts_with("#")) - .collect(); - - download_multiple(urls, outdir, file_count, conn_count, is_zippy).await?; - - } - - if let Some(url) = arguments.value_of("download") { - - download_multiple(vec![url.to_string()], outdir, 1, conn_count, is_zippy).await?; - - } - - if let Some(url) = arguments.value_of("zippy-resolve") { - - match zippy::resolve_link(&url).await { - Ok(resolved_url) => { - println!("{}", resolved_url); - }, - Err(_e) => { - println!("Zippyshare link could not be resolved"); + if !p_listfile.is_file() { + eprintln!("Listfile '{}' does not exist!", &listfile); exit(1); } + + let ifile = std::fs::File::open(p_listfile)?; + + cli_args.urls = std::io::BufReader::new(ifile) + .lines() + .map(|l| l.unwrap()) + .filter(|url| url.len() > 0 && !url.starts_with("#")) + .collect(); + }, + CLIAction::DownloadUrl(url) => { + cli_args.urls = vec![url.clone()]; } + CLIAction::ResolveZippyUrl(url) => { + match zippy::resolve_link(url).await { + Ok(resolved_url) => { + println!("{}", resolved_url); + }, + Err(_e) => { + println!("Zippyshare link could not be resolved"); + exit(1); + } + } + }, + + CLIAction::None => { + eprintln!("No action selected. This should not happen"); + exit(1); + } + } - Ok(()) + download_multiple(cli_args).await + } -async fn download_multiple(urls: Vec, outdir: &str, file_count: i32, conn_count: i32, is_zippy: bool) -> ResBE<()> { - let outdir = Path::new(outdir); +async fn download_multiple(cli_args: CLIArguments) -> ResBE<()> { + let outdir = cli_args.outdir; + let outdir = Path::new(&outdir); + + let file_count = cli_args.file_count; + + let zippy = cli_args.zippy; + + let conn_count = cli_args.conn_count; if !outdir.exists() { if let Err(_e) = std::fs::create_dir_all(outdir) { @@ -197,9 +270,9 @@ async fn download_multiple(urls: Vec, outdir: &str, file_count: i32, con let (tx, rx) = mpsc::unbounded_channel::(); - for offset in 0..file_count { + for offset in 0 .. file_count { - let urls: Vec = urls + let urls: Vec = cli_args.urls .iter() .enumerate() .filter(|(index, _)| (index) % file_count as usize == offset as usize) @@ -209,6 +282,8 @@ async fn download_multiple(urls: Vec, outdir: &str, file_count: i32, con let tx = tx.clone(); let outdir = outdir.to_owned(); let offset = offset; + let arg_filename = cli_args.into_file.clone(); + joiners.push(tokio::task::spawn(async move { for (i, url) in urls.iter().enumerate() { @@ -216,11 +291,11 @@ async fn download_multiple(urls: Vec, outdir: &str, file_count: i32, con let tx = tx.clone(); // Recalculated index in the main url vector, used as id - let global_url_index = i as i32 * file_count + offset; + let global_url_index = i as u32 * file_count + offset; let rep = DlReporter::new(global_url_index, tx); - let url = if is_zippy { + let url = if zippy { match zippy::resolve_link(&url).await { Ok(url) => url, Err(_e) => { @@ -235,7 +310,13 @@ async fn download_multiple(urls: Vec, outdir: &str, file_count: i32, con url.to_string() }; - let file_name = download::url_to_filename(&url); + let file_name = if let Some(arg_filename) = &arg_filename { + arg_filename.to_string() + } 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);