path: root/fatcat-cli/src/download.rs

use crate::{ApiModelIdent, Specifier};
use anyhow::{anyhow, Context, Result};
use crossbeam_channel as channel;
use fatcat_openapi::models::{FileEntity, ReleaseEntity};
use indicatif::{ProgressBar, ProgressStyle};
use log::{error, info};
use reqwest::header::USER_AGENT;
use sha1::Sha1;
use std::fmt;
use std::fs::File;
use std::io::{self, BufRead};
use std::path::PathBuf;
use std::thread;
use url::Url;

#[derive(Debug, PartialEq, Clone)]
pub enum DownloadStatus {
    Exists(String),
    Downloaded(String),
    NetworkError(String),
    HttpError(u16),
    PartialExists(String),
    NoPublicFile,
    FileMissingMetadata,
    WrongSize,
    WrongHash,
}

impl DownloadStatus {
    pub fn details(&self) -> Option<String> {
        match self {
            Self::Exists(p) => Some(p.to_string()),
            Self::Downloaded(p) => Some(p.to_string()),
            Self::NetworkError(p) => Some(p.to_string()),
            Self::PartialExists(p) => Some(p.to_string()),
            _ => None,
        }
    }
}

impl fmt::Display for DownloadStatus {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Exists(_) => write!(f, "exists"),
            Self::Downloaded(_) => write!(f, "success"),
            Self::NetworkError(_) => write!(f, "network-error"),
            Self::HttpError(p) => write!(f, "http-{}", p),
            Self::PartialExists(_) => write!(f, "partial-exists"),
            Self::NoPublicFile => write!(f, "no-public-file"),
            Self::FileMissingMetadata => write!(f, "missing-file-metadata"),
            Self::WrongSize => write!(f, "wrong-file-size"),
            Self::WrongHash => write!(f, "wrong-file-hash"),
        }
    }
}

struct Sha1WriteWrapper<W> {
    writer: W,
    hasher: Sha1,
}

impl<W> Sha1WriteWrapper<W> {
    fn new(writer: W) -> Self {
        Sha1WriteWrapper {
            writer,
            hasher: Sha1::new(),
        }
    }

    fn into_hexdigest(self) -> String {
        self.hasher.hexdigest()
    }
}

impl<W: std::io::Write> std::io::Write for Sha1WriteWrapper<W> {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        self.hasher.update(buf);
        self.writer.write(buf)
    }

    fn flush(&mut self) -> std::io::Result<()> {
        self.writer.flush()
    }
}

// eg, https://web.archive.org/web/20140802044207/http://www.geo.coop:80/sites/default/files/labs_of_oligarchy.pdf
fn rewrite_wayback_url(mut url: Url) -> Result<Url> {
    let mut segments: Vec<String> = url
        .path_segments()
        .unwrap()
        .map(|x| x.to_string())
        .collect();
    if segments.len() < 3 || url.host_str() != Some("web.archive.org") {
        return Err(anyhow!("not a valid wayback URL: {:?}", url));
    }
    // NOTE: the "12" digit timestamp here is to hack around a fraction of bad metadata in fatcat
    // catalog circa 2019. should be cleaned up eventually then this can be deprecated.
    if segments[0] == "web" && (segments[1].len() == 14 || segments[1].len() == 12) {
        segments[1] = format!("{}id_", segments[1]);
    }
    url.set_path(&segments.join("/"));
    Ok(url)
}

fn default_filename(specifier: &Specifier, fe: &FileEntity) -> PathBuf {
    let file_suffix = match fe.mimetype.as_deref() {
        Some("application/pdf") => ".pdf",
        Some("application/postscript") => ".ps",
        Some("text/html") => ".html",
        Some("text/xml") | Some("application/xml") | Some("application/xml+jats") => ".xml",
        // NOTE: most commonly .pdf if no type specified. should remove this default after updating
        // remaining file entities
        None => ".pdf",
        _ => "",
    };

    let path_string = format!("{}{}", specifier, file_suffix);
    PathBuf::from(&path_string)
}

/// Attempts to download a file entity, including verifying checksum.
pub fn download_file(
    fe: &FileEntity,
    specifier: &Specifier,
    output_path: Option<PathBuf>,
    show_progress: bool,
) -> Result<DownloadStatus> {
    let expected_sha1 = match &fe.sha1 {
        Some(v) => v,
        None => return Ok(DownloadStatus::FileMissingMetadata),
    };
    let expected_size = match fe.size {
        Some(v) => v as u64,
        None => return Ok(DownloadStatus::FileMissingMetadata),
    };

    let final_path = match output_path {
        Some(ref path) if path.is_dir() => {
            let mut full = output_path.unwrap_or_default();
            full.push(default_filename(specifier, fe));
            full
        }
        Some(path) => path,
        None => default_filename(specifier, fe),
    };

    // NOTE: this isn't perfect; there could have been a race condition
    if final_path.exists() {
        return Ok(DownloadStatus::Exists(
            final_path.to_string_lossy().to_string(),
        ));
    };

    let download_path = final_path.with_extension("partial_download");

    let raw_url = match fe.urls.as_ref() {
        None => return Ok(DownloadStatus::NoPublicFile),
        Some(url_list) if url_list.is_empty() => return Ok(DownloadStatus::NoPublicFile),
        // TODO: remove clone (?)
        Some(url_list) => {
            // prefer an IA_hosted URL, but fallback to any URL
            let public_url_list: Vec<&fatcat_openapi::models::FileUrl> = url_list
                .iter()
                .filter(|v| {
                    if let Ok(url) = Url::parse(&v.url) {
                        url.host_str() == Some("web.archive.org")
                            || url.host_str() == Some("archive.org")
                    } else {
                        false
                    }
                })
                .collect();
            if !public_url_list.is_empty() {
                public_url_list[0].url.clone()
            } else {
                url_list[0].url.clone()
            }
        }
    };

    let mut url = Url::parse(&raw_url)?;
    if url.host_str() == Some("web.archive.org") {
        url = rewrite_wayback_url(url)?;
    }

    let download_file = match std::fs::OpenOptions::new()
        .write(true)
        .create_new(true)
        .open(&download_path)
    {
        Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => {
            return Ok(DownloadStatus::PartialExists(
                download_path.to_string_lossy().to_string(),
            ))
        }
        Err(e) => return Err(e).context("opening temporary download file"),
        Ok(f) => f,
    };

    // TODO: print to stderr (?)
    info!("downloading: {}", url);
    let client = reqwest::blocking::Client::new();
    let mut resp = match client
        .get(url)
        .header(
            USER_AGENT,
            format!("fatcat-cli/{}", env!("CARGO_PKG_VERSION")),
        )
        .send()
    {
        Ok(r) => r,
        Err(e) => {
            std::fs::remove_file(download_path)?;
            return Ok(DownloadStatus::NetworkError(format!("{:?}", e)));
        }
    };

    if !resp.status().is_success() {
        std::fs::remove_file(download_path)?;
        return Ok(DownloadStatus::HttpError(resp.status().as_u16()));
    }

    if let Some(content_bytes) = resp.content_length() {
        if content_bytes != expected_size {
            std::fs::remove_file(download_path)?;
            return Ok(DownloadStatus::WrongSize);
        }
    }

    let (write_result, out_sha1) = match show_progress {
        true => {
            let pb = ProgressBar::new(fe.size.unwrap() as u64);
            pb.set_style(ProgressStyle::default_bar()
                .template("{spinner:.green} [{elapsed_precise}] [{bar:40}] {bytes}/{total_bytes} ({eta})")
                .progress_chars("#>-"));
            let mut wrapped_file = Sha1WriteWrapper::new(pb.wrap_write(download_file));
            let result = resp.copy_to(&mut wrapped_file);
            let out_sha1 = wrapped_file.into_hexdigest();
            (result, out_sha1)
        }
        false => {
            let mut wrapped_file = Sha1WriteWrapper::new(download_file);
            let result = resp.copy_to(&mut wrapped_file);
            let out_sha1 = wrapped_file.into_hexdigest();
            (result, out_sha1)
        }
    };

    let out_size = match write_result {
        Ok(r) => r,
        Err(e) => {
            std::fs::remove_file(download_path)?;
            return Ok(DownloadStatus::NetworkError(format!("{:?}", e)));
        }
    };

    if &out_sha1 != expected_sha1 {
        std::fs::remove_file(download_path)?;
        return Ok(DownloadStatus::WrongHash);
    }

    if out_size != expected_size {
        std::fs::remove_file(download_path)?;
        return Ok(DownloadStatus::WrongSize);
    }

    std::fs::rename(download_path, &final_path)?;
    Ok(DownloadStatus::Downloaded(
        final_path.to_string_lossy().to_string(),
    ))
}

pub fn download_release(
    re: &ReleaseEntity,
    output_path: Option<PathBuf>,
    show_progress: bool,
) -> Result<DownloadStatus> {
    let file_entities = match &re.files {
        None => {
            return Err(anyhow!(
                "expected file sub-entities to be 'expanded' on release"
            ))
        }
        Some(list) => list,
    };
    let mut status = DownloadStatus::NoPublicFile;
    let specifier = re.specifier();
    for fe in file_entities {
        status = download_file(&fe, &specifier, output_path.clone(), show_progress)?;
        match status {
            DownloadStatus::Exists(_) | DownloadStatus::Downloaded(_) => break,
            _ => (),
        };
    }
    Ok(status)
}

/// Tries either file or release
fn download_entity(
    json_str: String,
    output_path: Option<PathBuf>,
    show_progress: bool,
) -> Result<(DownloadStatus, String)> {
    let release_attempt = serde_json::from_str::<ReleaseEntity>(&json_str);
    if let Ok(re) = release_attempt {
        if re.ident.is_some() && (re.title.is_some() || re.files.is_some()) {
            let status = download_release(&re, output_path, show_progress)?;
            let status_line = format!(
                "release_{}\t{}\t{}",
                re.ident.unwrap(),
                status,
                status.details().unwrap_or_else(|| "".to_string())
            );
            return Ok((status, status_line));
        };
    }
    let file_attempt =
        serde_json::from_str::<FileEntity>(&json_str).context("parsing entity for download");
    match file_attempt {
        Ok(fe) => {
            if fe.ident.is_some() && fe.urls.is_some() {
                let specifier = fe.specifier();
                let status = download_file(&fe, &specifier, output_path, show_progress)?;
                let status_line = format!(
                    "file_{}\t{}\t{}",
                    fe.ident.unwrap(),
                    status,
                    status.details().unwrap_or_else(|| "".to_string())
                );
                Ok((status, status_line))
            } else {
                Err(anyhow!("not a file entity (JSON)"))
            }
        }
        Err(e) => Err(e),
    }
}

struct DownloadTask {
    json_str: String,
    output_path: Option<PathBuf>,
    show_progress: bool,
}

fn loop_printer(
    output_receiver: channel::Receiver<String>,
    done_sender: channel::Sender<()>,
) -> Result<()> {
    for line in output_receiver {
        println!("{}", line);
    }
    done_sender.send(())?;
    drop(done_sender);
    Ok(())
}

fn loop_download_tasks(
    task_receiver: channel::Receiver<DownloadTask>,
    output_sender: channel::Sender<String>,
) {
    let thread_result: Result<()> = (|| {
        for task in task_receiver {
            let (_, status_line) =
                download_entity(task.json_str, task.output_path, task.show_progress)?;
            output_sender.send(status_line)?;
        }
        drop(output_sender);
        Ok(())
    })();
    if let Err(ref e) = thread_result {
        error!("{}", e);
    }
    thread_result.unwrap()
}

pub fn download_batch(
    input_path: Option<PathBuf>,
    output_dir: Option<PathBuf>,
    limit: Option<u64>,
    jobs: u64,
) -> Result<u64> {
    let mut count = 0;

    assert!(jobs > 0 && jobs <= 12);

    let show_progress = jobs == 1;

    let (task_sender, task_receiver) = channel::bounded(12);
    let (output_sender, output_receiver) = channel::bounded(12);
    let (done_sender, done_receiver) = channel::bounded(0);

    for _ in 0..jobs {
        let task_receiver = task_receiver.clone();
        let output_sender = output_sender.clone();
        thread::spawn(move || {
            loop_download_tasks(task_receiver, output_sender);
        });
    }
    drop(output_sender);

    // Start printer thread
    thread::spawn(move || loop_printer(output_receiver, done_sender).expect("printing to stdout"));

    match input_path {
        None => {
            let stdin = io::stdin();
            let stdin_lock = stdin.lock();
            let lines = stdin_lock.lines();
            for line in lines {
                let task = DownloadTask {
                    json_str: line?,
                    output_path: output_dir.clone(),
                    show_progress,
                };
                task_sender.send(task)?;
                count += 1;
                if let Some(limit) = limit {
                    if count >= limit {
                        break;
                    }
                }
            }
        }
        Some(path) => {
            let input_file = File::open(path)?;
            let buffered = io::BufReader::new(input_file);
            let lines = buffered.lines();
            for line in lines {
                let task = DownloadTask {
                    json_str: line?,
                    output_path: output_dir.clone(),
                    show_progress,
                };
                task_sender.send(task)?;
                count += 1;
                if let Some(limit) = limit {
                    if count >= limit {
                        break;
                    }
                }
            }
        }
    };
    drop(task_sender);
    done_receiver.recv()?;
    Ok(count)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_rewrite_wayback_url() -> () {
        assert_eq!(rewrite_wayback_url(Url::parse("https://web.archive.org/web/20140802044207/http://www.geo.coop:80/sites/default/files/labs_of_oligarchy.pdf").unwrap()).unwrap(),
        Url::parse("https://web.archive.org/web/20140802044207id_/http://www.geo.coop:80/sites/default/files/labs_of_oligarchy.pdf").unwrap());

        // not a wayback URL
        assert!(rewrite_wayback_url(Url::parse("https://fatcat.wiki/blah.pdf").unwrap()).is_err());

        // too short
        assert!(
            rewrite_wayback_url(Url::parse("https://web.archive.org/file.pdf").unwrap()).is_err()
        );
    }
}