From 4e79ab051917fece21fcec0bd31cc2942462fda0 Mon Sep 17 00:00:00 2001 From: Melody Horn Date: Sat, 27 Mar 2021 09:02:59 -0600 Subject: rough draft --- src/download.rs | 46 ++++++++++++++++++++++++++++++++++++++++++++++ src/main.rs | 39 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 85 insertions(+) create mode 100644 src/download.rs create mode 100644 src/main.rs (limited to 'src') diff --git a/src/download.rs b/src/download.rs new file mode 100644 index 0000000..537dae3 --- /dev/null +++ b/src/download.rs @@ -0,0 +1,46 @@ +use std::path::{Path, PathBuf}; + +use anyhow::Error; +use fehler::throws; +use hyper::{body::HttpBody, Client, header::CONTENT_TYPE, Uri}; +use tokio::fs; +use tokio::io::AsyncWriteExt; + +use super::Args; + +#[throws] +pub(crate) async fn download(url: Uri, args: &Args) { + let output_file_path = if let Some(output) = &args.output_document { + Some(output.clone()) + } else { + let url_path = Path::new(url.path()); + match (url_path.file_name(), url.path().ends_with('/')) { + (Some(file_name), false) => Some(PathBuf::from(file_name)), + _ => None, + } + }; + + let client = Client::new(); + let mut response = client.get(url.clone()).await?; + if !response.status().is_success() { + panic!("non-success response code {} in URL {}", response.status(), url); + } + + let output_file_path = if let Some(path) = output_file_path { + path + } else { + let content_type = response.headers().get(CONTENT_TYPE); + let extension = content_type + .and_then(|mime_type| mime_type.to_str().ok()) + .and_then(|mime_type| mime2ext::mime2ext(mime_type)) + .map(|x| format!(".{}", x)) + .unwrap_or_default(); + PathBuf::from(format!("index{}", extension)) + }; + + let mut output_file = fs::File::create(output_file_path).await.expect("couldn't open output file!"); + + while let Some(data) = response.body_mut().data().await { + output_file.write_all(&data?).await?; + } +} diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..da1ae1f --- /dev/null +++ b/src/main.rs @@ -0,0 +1,39 @@ +use std::mem; +use std::path::PathBuf; + +use anyhow::Error; +use hyper::Uri; +use structopt::StructOpt; + +mod download; +use download::download; + +#[derive(StructOpt, Debug)] +struct Args { + /// Specify output filename (will cause an error if multiple URLs are given) + #[structopt(short = "O", long)] + output_document: Option, + + /// The URLs to download + urls: Vec, +} + +#[tokio::main] +async fn main() -> Result<(), Error> { + let (args, urls) = { + let mut args: Args = Args::from_args(); + let urls = mem::take(&mut args.urls); + (args, urls) + }; + if urls.len() != 1 && args.output_document.is_some() { + panic!("can't use the same output file for multiple URLs!") + } + + let download_handles = urls + .into_iter() + .map(|url| download(url, &args)); + let downloads = futures::future::join_all(download_handles); + let _downloads = downloads.await; + + Ok(()) +} -- cgit v1.2.3