use std::{ collections::VecDeque, io::BufRead, path::Path, process::exit, sync::Arc, sync::Mutex, time::SystemTime, }; use clap::Parser; use futures::future::join_all; use tokio::sync::mpsc; use crate::{ args::{CLIAction, CLIArgs}, dlreport::{DlReport, DlReporter, DlStatus}, errors::ResBE, }; mod args; mod dlreport; mod download; mod errors; mod zippy; #[tokio::main] 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, }; 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) .lines() .map(|l| l.unwrap()) .filter(|url| url.len() > 0 && !url.starts_with("#")) .collect() } async fn download_multiple(cli_args: CLIArgs, raw_urls: Vec) -> ResBE<()> { let outdir = cli_args.outdir; let outdir = Path::new(&outdir); 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())); raw_urls.iter().enumerate().for_each(|(i, url)| { urls.lock().unwrap().push_back((i as u32, url.clone())); }); if !outdir.exists() { if let Err(_e) = std::fs::create_dir_all(outdir) { eprintln!( "Error creating output directory '{}'", outdir.to_str().unwrap() ); 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 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(), }); } }; } })) } drop(tx); dlreport::watch_and_print_reports(rx, raw_urls.len() as i32).await?; join_all(joiners).await; println!("Total time: {}s", t_start.elapsed()?.as_secs()); Ok(()) }