use std::{
fs::{self, File},
future::Future,
io::{Read, Write},
path::Path,
pin::Pin,
sync::Arc,
};
use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
use sha1::{Digest, Sha1};
use sha2::{Sha224, Sha256, Sha384, Sha512};
use smol::{Executor, io::AsyncReadExt, lock::Semaphore};
use surf::Client;
#[cfg(feature = "cas")]
pub mod cas;
#[cfg(feature = "decompress")]
pub mod decompress;
mod redirection_middleware;
pub struct Downloader {
pub files: Vec<DLFile>,
pub max_concurrent_downloads: usize,
pub max_redirections: usize,
pub style: ProgressStyle,
}
pub struct DLFile {
pub size: u64,
pub url: String,
pub hashes: DLHashes,
pub path: String,
#[cfg(feature = "decompress")]
pub decompression_config: Option<decompress::DLDecompressionConfig>,
pub on_download: Arc<dyn Fn(String) + Send + Sync>,
#[cfg(feature = "cas")]
pub cas: Option<cas::DLStorage>,
}
#[derive(Debug, Clone)]
pub struct DLHashes {
pub hashes: Vec<(DLHashType, String)>,
}
impl DLHashes {
pub fn new() -> Self {
Self { hashes: Vec::new() }
}
pub fn add_hash(mut self, hash_type: DLHashType, hash_value: String) -> Self {
self.hashes.push((hash_type, hash_value));
self
}
pub fn sha1(mut self, hash: &str) -> Self {
self.hashes.push((DLHashType::SHA1, hash.to_string()));
self
}
pub fn sha256(mut self, hash: &str) -> Self {
self.hashes.push((DLHashType::SHA256, hash.to_string()));
self
}
pub fn sha384(mut self, hash: &str) -> Self {
self.hashes.push((DLHashType::SHA384, hash.to_string()));
self
}
pub fn sha512(mut self, hash: &str) -> Self {
self.hashes.push((DLHashType::SHA512, hash.to_string()));
self
}
pub fn sha224(mut self, hash: &str) -> Self {
self.hashes.push((DLHashType::SHA224, hash.to_string()));
self
}
pub fn verify_data(&self, data: &[u8]) -> bool {
self.hashes
.iter()
.find(|hashed| {
let (typ, hash) = hashed;
typ.verify_data(data, hash)
})
.is_some()
}
pub fn verify_str(&self, data: &str) -> bool {
self.verify_data(data.as_bytes())
}
pub fn verify_file(&self, path: &str) -> bool {
let data = std::fs::read(path).unwrap();
self.verify_data(&data)
}
}
#[derive(Debug, Clone)]
pub enum DLHashType {
SHA1,
SHA256,
SHA224,
SHA384,
SHA512,
}
impl DLHashType {
fn compute_hash<D: Digest + Default>(data: &[u8]) -> String {
let mut hasher = D::default();
hasher.update(data);
let result = hasher.finalize();
hex::encode(result)
}
pub fn compute(&self, data: &[u8]) -> String {
match self {
DLHashType::SHA1 => Self::compute_hash::<Sha1>(data),
DLHashType::SHA256 => Self::compute_hash::<Sha256>(data),
DLHashType::SHA224 => Self::compute_hash::<Sha224>(data),
DLHashType::SHA384 => Self::compute_hash::<Sha384>(data),
DLHashType::SHA512 => Self::compute_hash::<Sha512>(data),
}
}
pub fn verify_str(&self, data: &str, hash: &str) -> bool {
self.compute(data.to_string().as_bytes()) == hash
}
pub fn verify_data(&self, data: &[u8], hash: &str) -> bool {
self.compute(data) == hash
}
pub fn verify_file(&self, path: &Path, hash: &str) -> bool {
let mut buffer = Vec::new();
File::open(path)
.expect("Failed to open file")
.read_to_end(&mut buffer)
.expect("Failed to read file");
self.verify_data(buffer.as_slice(), hash)
}
}
fn symlink_exists(path: &Path) -> bool {
match fs::symlink_metadata(path) {
Ok(metadata) => metadata.file_type().is_symlink(),
Err(_) => false,
}
}
impl DLFile {
pub async fn download(&self, progress: ProgressBar, client: Client) -> Result<(), String> {
let url = self.url.clone();
let path = self.path.clone();
let hashes = self.hashes.clone();
let size = self.size;
let path_clone = self.path.clone();
let mut response = client.get(&url).await.expect("Failed to get response");
progress.set_length(size);
progress.set_message(format!("Downloading {}", path_clone));
let path_hash: String = if response.status().is_success() {
let ppath = Path::new(&path);
if let Some(parent) = ppath.parent() {
fs::create_dir_all(parent).unwrap();
}
let (mut file, path_hash) = if self.cas.is_some() && hashes.hashes.len() > 0 {
let hash = hashes.hashes.get(0).unwrap().clone().1;
let storage = self.cas.as_ref().unwrap();
if storage.find(hash.as_str()).is_some() {
if !symlink_exists(Path::new(path.clone().as_str())) {
storage.symlink(hash.as_str(), path.clone().as_str());
}
progress.set_message(format!("Done {}", path_clone));
progress.set_position(size);
return Ok(());
}
(
storage.new_file(hash.as_str(), path.clone().as_str()),
storage.path(hash.as_str()),
)
} else {
(File::create(path.clone()).unwrap(), path.clone())
};
let mut downloaded = 0;
let mut buffer = [0; 8192];
let mut body = response.take_body();
loop {
match AsyncReadExt::read(&mut body, &mut buffer).await {
Ok(0) => break, Ok(n) => {
file.write_all(&buffer[..n]).unwrap();
downloaded += n as u64;
progress.set_position(downloaded);
}
Err(e) => return Err(e.to_string()),
}
}
path_hash
} else {
progress.abandon_with_message(format!("Error: {}", response.status()));
String::new()
};
if hashes.hashes.len() > 0 && !hashes.verify_file(&path_hash) {
progress.abandon_with_message(format!("Hash verification failed for {}", path_clone));
return Err("Hash verification failed".to_string());
}
(self.on_download)(path_clone.clone());
#[cfg(feature = "decompress")]
{
if self.decompression_config.is_some() {
progress.set_message("Decompressing...");
let config = self.decompression_config.as_ref().unwrap();
config.decompress(&path_clone)?;
if config.delete_after {
progress.set_message("Cleaning up...");
std::fs::remove_file(&path_clone).expect("Failed to delete file");
}
}
}
progress.finish_with_message(format!("DONE {}", path));
Ok(())
}
pub fn new() -> Self {
DLFile {
path: String::new(),
url: String::new(),
size: 0,
hashes: DLHashes::new(),
#[cfg(feature = "decompress")]
decompression_config: None,
on_download: Arc::new(|_| {}),
cas: None,
}
}
pub fn with_path(mut self, path: &str) -> Self {
self.path = path.to_string();
self
}
pub fn with_url(mut self, url: &str) -> Self {
self.url = url.to_string();
self
}
pub fn with_size(mut self, size: u64) -> Self {
self.size = size;
self
}
pub fn with_hashes(mut self, hashes: DLHashes) -> Self {
self.hashes = hashes;
self
}
#[cfg(feature = "decompress")]
pub fn with_decompression_config(
mut self,
config: crate::decompress::DLDecompressionConfig,
) -> Self {
self.decompression_config = Some(config);
self
}
pub fn with_on_download(mut self, on_download: Arc<dyn Fn(String) + Send + Sync>) -> Self {
self.on_download = on_download;
self
}
#[cfg(feature = "cas")]
pub fn with_cas(mut self, value: cas::DLStorage) -> Self {
self.cas = Some(value);
self
}
}
impl Downloader {
pub fn new() -> Self {
Downloader {
files: Vec::new(),
max_concurrent_downloads: 5,
max_redirections: 5,
style: ProgressStyle::with_template(
"[{elapsed_precise}] {bar:40.green/red} {pos:>7}/{len:7} {msg}",
)
.unwrap()
.progress_chars("##-"),
}
}
pub fn with_files(mut self, files: Vec<DLFile>) -> Self {
self.files.extend(files);
self
}
pub fn from_files(files: Vec<DLFile>) -> Self {
let builder = Downloader::new();
builder.with_files(files)
}
pub fn add_file(mut self, file: DLFile) -> Self {
self.files.push(file);
self
}
pub fn start(&self) {
let m = MultiProgress::new();
let semaphore = Arc::new(Semaphore::new(self.max_concurrent_downloads));
let executor = Arc::new(Executor::new());
let client = Client::new().with(redirection_middleware::RedirectMiddleware::new(
self.max_redirections,
));
let futures: Vec<Pin<Box<dyn Future<Output = Result<(), String>>>>> = self
.files
.iter()
.map(|dl_file| {
let progress = m.add(ProgressBar::new(0).with_style(self.style.clone()));
let semaphore = Arc::clone(&semaphore);
let client = client.clone();
let task: Pin<Box<dyn Future<Output = Result<(), String>>>> =
Box::pin(executor.run(async move {
let permit = semaphore.acquire().await;
#[cfg(feature = "no_static_client")]
let client = create_client(self.max_redirections);
dl_file.download(progress, client.clone()).await?;
drop(permit);
Ok(())
}));
task
})
.collect();
smol::block_on(async {
futures::future::join_all(futures).await;
});
}
pub fn with_style(mut self, style: ProgressStyle) -> Self {
self.style = style;
self
}
pub fn with_max_concurrent_downloads(mut self, max_concurrent_downloads: usize) -> Self {
self.max_concurrent_downloads = max_concurrent_downloads;
self
}
pub fn with_max_redirections(mut self, max_redirections: usize) -> Self {
self.max_redirections = max_redirections;
self
}
}
fn create_client(max_redirections: usize) -> Client {
Client::new().with(redirection_middleware::RedirectMiddleware::new(
max_redirections,
))
}