126 lines
3.6 KiB
Rust
126 lines
3.6 KiB
Rust
use std::{
|
|
error, fmt,
|
|
fs::{File, OpenOptions},
|
|
io::{self, Read, Write},
|
|
};
|
|
|
|
use bytes::Bytes;
|
|
use futures_util::StreamExt;
|
|
use indicatif::{ProgressBar, ProgressState, ProgressStyle};
|
|
use reqwest::{self, Client};
|
|
use std::cmp;
|
|
|
|
use crate::get_mut_ref;
|
|
|
|
pub enum DualWriter {
|
|
File(File),
|
|
Buffer(Vec<u8>),
|
|
}
|
|
|
|
impl DualWriter {
|
|
pub fn write(&mut self, bytes: Bytes) -> Result<(), std::io::Error> {
|
|
match self {
|
|
Self::Buffer(buffer) => {
|
|
bytes.into_iter().for_each(|byte| buffer.push(byte));
|
|
}
|
|
Self::File(file) => {
|
|
file.write(&bytes)?;
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
pub fn len(&self) -> u64 {
|
|
match self {
|
|
DualWriter::File(f) => f.metadata().and_then(|x| Ok(x.len())).unwrap_or_else(|e| {
|
|
println!("Could not get metadata from file {:?}", e);
|
|
0
|
|
}),
|
|
DualWriter::Buffer(buf) => buf.len() as u64,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl TryFrom<Option<&str>> for DualWriter {
|
|
type Error = io::Error;
|
|
|
|
fn try_from(file_name: Option<&str>) -> Result<Self, Self::Error> {
|
|
Ok(if let Some(file_name) = file_name {
|
|
let file = match OpenOptions::new().append(true).open(&file_name) {
|
|
Ok(f) => f,
|
|
Err(e) => {
|
|
let e = e as io::Error;
|
|
if e.kind() == io::ErrorKind::NotFound {
|
|
File::create(&file_name)?
|
|
} else {
|
|
return Err(e);
|
|
}
|
|
}
|
|
};
|
|
Self::File(file)
|
|
} else {
|
|
Self::Buffer(Vec::<u8>::new())
|
|
})
|
|
}
|
|
}
|
|
|
|
impl TryInto<String> for DualWriter {
|
|
type Error = String;
|
|
|
|
fn try_into(self) -> Result<String, Self::Error> {
|
|
Ok(match self {
|
|
Self::Buffer(buffer) => {
|
|
String::from_utf8(buffer).or(Err("Failed to decode buffer".to_owned()))?
|
|
}
|
|
Self::File(file) => {
|
|
let mut buf = String::new();
|
|
|
|
// Well this is safe since I consume the file anyways
|
|
let file = unsafe { get_mut_ref(&file) };
|
|
file.read_to_string(&mut buf)
|
|
.or(Err("Failed to read file".to_owned()))?;
|
|
buf
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
pub async fn download_with_progress(
|
|
link: &str,
|
|
file_name: Option<&str>,
|
|
) -> Result<DualWriter, Box<dyn error::Error>> {
|
|
let mut dual_writer: DualWriter = file_name.try_into()?;
|
|
|
|
let client = Client::builder().gzip(true).deflate(true).build()?;
|
|
let builder = client
|
|
.get(link)
|
|
.header("Range", format!("bytes={}-", dual_writer.len()));
|
|
|
|
let resp = builder.send().await?;
|
|
|
|
let content_length = resp.content_length().unwrap_or_default();
|
|
if content_length == 0 {
|
|
println!("File was already downloaded");
|
|
return Ok(dual_writer);
|
|
}
|
|
|
|
let progress_bar = ProgressBar::new(content_length);
|
|
progress_bar.set_style(ProgressStyle::with_template("{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({eta})").unwrap()
|
|
.with_key("eta", |state: &ProgressState, w: &mut dyn fmt::Write| write!(w, "{:.1}m", state.eta().as_secs_f64() / 60.0 ).unwrap())
|
|
.progress_chars("#>-"));
|
|
|
|
let mut downloaded: u64 = 0;
|
|
let mut stream = resp.bytes_stream();
|
|
|
|
while let Some(item) = stream.next().await {
|
|
let bytes = item.unwrap();
|
|
downloaded = cmp::min(downloaded + (bytes.len() as u64), content_length);
|
|
dual_writer.write(bytes)?;
|
|
|
|
progress_bar.set_position(downloaded);
|
|
}
|
|
|
|
Ok(dual_writer)
|
|
}
|