ffdl/src/main.rs

use std::path::Path;
use std::process::exit;
use clap::{ App, Arg, ArgGroup, crate_version };
use tokio::sync::mpsc;
use futures::future::join_all;
use std::io::BufRead;
use std::time::SystemTime;

use dlreport::{ DlReport, DlStatus, DlReporter };
use errors::ResBE;

mod zippy;
mod download;
mod errors;
mod dlreport;


#[tokio::main]
async fn main() -> ResBE<()> {

    let arguments = App::new("FFDL - Fast File Downloader")
        .version(crate_version!())
        .about("Download files fast")
        .arg(
            Arg::with_name("outdir")
                .short("o")
                .long("outdir")
                .value_name("OUTPUT DIR")
                .takes_value(true)
                .help("Set the output directory. The directory will be created \
                        if it doesn't exit yet")
        )
        .arg(
            Arg::with_name("file_count")
                .short("n")
                .long("num-files")
                .value_name("NUMBER OF CONCURRENT FILE DOWNLOADS")
                .takes_value(true)
                .help("Specify the number concurrent file downloads")
        )
        .arg(
            Arg::with_name("conn_count")
                .short("c")
                .long("connections")
                .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")
        )
        .arg(
            Arg::with_name("zippyshare")
                .short("z")
                .long("zippy")
                .takes_value(false)
                .help("The provided URLs are zippyshare URLs and need to be resolved")
        )
        .group(
            ArgGroup::with_name("action")
                .required(true)
        )
        .arg(
            Arg::with_name("listfile")
                .short("l")
                .long("listfile")
                .value_name("URL LIST")
                .takes_value(true)
                .group("action")
                .help("Download all files form the specified url list")
        )
        .arg(
            Arg::with_name("download")
                .short("d")
                .long("download")
                .value_name("URL")
                .takes_value(true)
                .group("action")
                .help("Download only the specified URL")
        )
        .arg(
            Arg::with_name("zippy-resolve")
                .long("zippy-resolve")
                .value_name("ZIPPYSHARE URL")
                .takes_value(true)
                .group("action")
                .help("Resolve the zippyshare url to real download url")
        )
        .get_matches();
    

    let outdir = match arguments.value_of("outdir") {
        Some(it) => it,
        None => "./"
    };

    let file_count = match arguments.value_of("file_count") {
        Some(it) => it,
        None => "1"
    };

    let file_count: i32 = match file_count.parse() {
        Ok(it) => it,
        Err(_) => {
            eprintln!("Invalid value for num-files: {}", file_count);
            exit(1);
        }
    };
    
    if file_count <= 0 {
        eprintln!("Invalid value for num-files: {}", file_count);
        exit(1);
    }

    let conn_count = match arguments.value_of("conn_count") {
        Some(it) => it,
        None => "1"
    };
    
    let conn_count: i32 = match conn_count.parse() {
        Ok(it) => it,
        Err(_) => {
            eprintln!("Invalid value for connections: {}", conn_count);
            exit(1);
        }
    };
    
    if conn_count <= 0 {
        eprintln!("Invalid value for connections: {}", conn_count);
        exit(1);
    }

    let is_zippy = arguments.is_present("zippyshare");


    // 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") {

        let listfile = Path::new(s_listfile);

        if !listfile.is_file() {
            eprintln!("Listfile '{}' does not exist!", s_listfile);
            exit(1);
        }

        let ifile = std::fs::File::open(listfile)?;

        let urls: Vec<String> = 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");
                exit(1);
            }

        }

    }

    Ok(())
}

async fn download_multiple(urls: Vec<String>, outdir: &str, file_count: i32, conn_count: i32, is_zippy: bool) -> ResBE<()> {
    let outdir = Path::new(outdir);

    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::<DlReport>();

    for offset in 0..file_count {

        let urls: Vec<String> = urls
            .iter()
            .enumerate()
            .filter(|(index, _)| (index) % file_count as usize == offset as usize)
            .map(|(_, v)| v.to_owned())
            .collect();
        
        let tx = tx.clone();
        let outdir = outdir.to_owned();
        let offset = offset;
        joiners.push(tokio::task::spawn(async move {

            for (i, url) in urls.iter().enumerate() {

                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 rep = DlReporter::new(global_url_index, tx);

                let url = if is_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 = 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);
            
                // If file with same name is present locally, check filesize
                if path_into_file.exists() {
                    let filesize = match download::http_get_filesize_and_range_support(&url).await {
                        Ok((filesize, _)) => filesize,
                        Err(_e) => {
                            rep.send(
                                DlStatus::Message(format!("Error while querying metadata: {}", url))
                            );
                            continue;
                        }
                    };

                    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)));
                        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()).await {
                        rep.send(DlStatus::DoneErr {
                            filename: into_file.to_string()
                        });
                    }
                } else {
                    if let Err(_e) = download::download_feedback_multi(&url, &into_file, rep.clone(), conn_count).await {
                        rep.send(DlStatus::DoneErr {
                            filename: into_file.to_string()
                        });
                    }
                };

            }
        }))

    }
    
    drop(tx);

    dlreport::watch_and_print_reports(rx).await?;

    join_all(joiners).await;

    println!("Total time: {}s", t_start.elapsed()?.as_secs());

    Ok(())
}