use crate::DownloadOptions;
use crate::config::ProxyConfig;
use crate::optimization::Optimizer;
use crate::progress::create_progress_bar;
use crate::utils::{self, print};
use humansize::{DECIMAL, format_size};
use mime::Mime;
use reqwest::blocking::Client;
use reqwest::header::{CONTENT_DISPOSITION, CONTENT_LENGTH, CONTENT_TYPE};
use sha2::Digest;
use std::error::Error;
use std::fs::File;
use std::io::{Read, Write};
use std::path::{Path, PathBuf};
use std::time::{Duration, Instant};
const MAX_RETRIES: u32 = 3;
const RETRY_DELAY: Duration = Duration::from_secs(2);
pub fn check_disk_space(
path: &Path,
required_size: u64,
) -> Result<(), Box<dyn Error + Send + Sync>> {
let dir = path.parent().unwrap_or(Path::new("."));
let available_space = fs2::available_space(dir)?;
if available_space < required_size {
return Err(format!(
"Insufficient disk space. Required: {}, Available: {}",
format_size(required_size, DECIMAL),
format_size(available_space, DECIMAL)
)
.into());
}
Ok(())
}
pub fn validate_filename(filename: &str) -> Result<(), Box<dyn Error + Send + Sync>> {
if filename.is_empty() {
return Err("Filename cannot be empty".into());
}
if filename.contains('/') || filename.contains('\\') {
return Err("Filename cannot contain directory separators".into());
}
if filename.contains('\0') {
return Err("Filename cannot contain null bytes".into());
}
if filename.contains("..") {
return Err("Filename cannot contain path traversal sequences".into());
}
if filename.len() > 255 {
return Err("Filename exceeds maximum length of 255 bytes".into());
}
let stem = filename.split('.').next().unwrap_or(filename);
const RESERVED: &[&str] = &[
"CON", "PRN", "AUX", "NUL",
"COM1", "COM2", "COM3", "COM4", "COM5", "COM6", "COM7", "COM8", "COM9",
"LPT1", "LPT2", "LPT3", "LPT4", "LPT5", "LPT6", "LPT7", "LPT8", "LPT9",
];
if RESERVED.iter().any(|r| stem.eq_ignore_ascii_case(r)) {
return Err(format!("'{}' is a reserved filename on Windows", stem).into());
}
Ok(())
}
pub fn download(
target: &str,
proxy: ProxyConfig,
optimizer: Optimizer,
options: DownloadOptions,
status_callback: Option<&(dyn Fn(String) + Send + Sync)>,
) -> Result<(), Box<dyn Error + Send + Sync>> {
let quiet_mode = options.quiet_mode;
let mut client_builder = Client::builder()
.timeout(Duration::from_secs(30))
.user_agent(concat!("KGet/", env!("CARGO_PKG_VERSION")))
.no_gzip()
.no_deflate();
if proxy.enabled {
if let Some(proxy_url) = &proxy.url {
let proxy_client = match proxy.proxy_type {
crate::config::ProxyType::Http => reqwest::Proxy::http(proxy_url),
crate::config::ProxyType::Https => reqwest::Proxy::https(proxy_url),
crate::config::ProxyType::Socks5 => reqwest::Proxy::all(proxy_url),
};
if let Ok(mut proxy_client) = proxy_client {
if let (Some(username), Some(password)) = (&proxy.username, &proxy.password) {
proxy_client = proxy_client.basic_auth(username, password);
}
client_builder = client_builder.proxy(proxy_client);
}
}
}
let client = client_builder.build()?;
let mut retries = 0;
let response = loop {
let mut req = client.get(target);
for (name, value) in &options.extra_headers {
if let (Ok(n), Ok(v)) = (
reqwest::header::HeaderName::from_bytes(name.as_bytes()),
reqwest::header::HeaderValue::from_str(value),
) {
req = req.header(n, v);
}
}
match req.send() {
Ok(resp) => {
let status = resp.status();
if status.is_success() {
break resp;
} else if status.is_server_error() {
retries += 1;
if retries >= MAX_RETRIES {
return Err(
format!("HTTP {} after {} attempts", status, MAX_RETRIES).into()
);
}
print(
&format!(
"Server error {} on attempt {}, retrying in {} seconds...",
status,
retries,
RETRY_DELAY.as_secs()
),
quiet_mode,
);
std::thread::sleep(RETRY_DELAY);
} else {
return Err(format!("HTTP error: {}", status).into());
}
}
Err(e) => {
retries += 1;
if retries >= MAX_RETRIES {
return Err(format!("Failed after {} attempts: {}", MAX_RETRIES, e).into());
}
print(
&format!(
"Attempt {} failed, retrying in {} seconds...",
retries,
RETRY_DELAY.as_secs()
),
quiet_mode,
);
std::thread::sleep(RETRY_DELAY);
}
}
};
print(
&format!("HTTP request sent... {}", response.status()),
quiet_mode,
);
let content_length = response
.headers()
.get(CONTENT_LENGTH)
.and_then(|ct_len| ct_len.to_str().ok())
.and_then(|s| s.parse::<u64>().ok());
let content_type = response
.headers()
.get(CONTENT_TYPE)
.and_then(|ct| ct.to_str().ok())
.and_then(|s| s.parse::<Mime>().ok());
let server_filename = response
.headers()
.get(CONTENT_DISPOSITION)
.and_then(|v| v.to_str().ok())
.and_then(parse_content_disposition_filename);
if let Some(len) = content_length {
print(
&format!("Length: {} ({})", len, format_size(len, DECIMAL)),
quiet_mode,
);
} else {
print("Length: unknown", quiet_mode);
}
if let Some(ref ct) = content_type {
print(&format!("Type: {}", ct), quiet_mode);
}
let is_iso = target.to_lowercase().ends_with(".iso")
|| content_type.as_ref().map_or(false, |ct| {
ct.essence_str() == "application/x-iso9660-image"
|| ct.essence_str() == "application/x-cd-image"
});
if is_iso {
print(
"ISO file detected. Ensuring raw download to prevent corruption...",
quiet_mode,
);
}
let tentative_path: PathBuf;
if let Some(output_arg_str) = options.output_path {
let user_path = PathBuf::from(output_arg_str.clone());
let is_target_dir =
user_path.is_dir() || output_arg_str.ends_with(std::path::MAIN_SEPARATOR);
if is_target_dir {
let base_filename = utils::get_filename_from_url_or_default(target, "downloaded_file");
validate_filename(&base_filename)?;
tentative_path = user_path.join(base_filename);
} else {
if let Some(file_name_osstr) = user_path.file_name() {
if let Some(file_name_str) = file_name_osstr.to_str() {
if file_name_str.is_empty() {
return Err(format!(
"Invalid output path, does not specify a file name: {}",
user_path.display()
)
.into());
}
validate_filename(file_name_str)?;
} else {
return Err("Output filename contains invalid characters (non-UTF-8)".into());
}
} else {
return Err(format!(
"Invalid output path, does not specify a file name: {}",
user_path.display()
)
.into());
}
tentative_path = user_path;
}
} else {
let base_filename = if let Some(ref name) = server_filename {
name.clone()
} else {
utils::get_filename_from_url_or_default(target, "downloaded_file")
};
validate_filename(&base_filename)?;
tentative_path = PathBuf::from(base_filename);
}
let final_path: PathBuf = if tentative_path.is_absolute() {
tentative_path
} else {
let current_dir = std::env::current_dir()
.map_err(|e| format!("Failed to get current directory: {}", e))?;
current_dir.join(tentative_path)
};
if let Some(parent_dir) = final_path.parent() {
if !parent_dir.as_os_str().is_empty()
&& parent_dir != Path::new("/")
&& !parent_dir.exists()
{
std::fs::create_dir_all(parent_dir).map_err(|e| {
format!("Failed to create directory {}: {}", parent_dir.display(), e)
})?;
if !quiet_mode {
print(
&format!("Created directory: {}", parent_dir.display()),
quiet_mode,
);
}
}
}
if !quiet_mode {
print(&format!("Saving to: {}", final_path.display()), quiet_mode);
}
if let Some(len) = content_length {
check_disk_space(&final_path, len)?;
}
let mut dest = File::create(&final_path)
.map_err(|e| format!("Failed to create file {}: {}", final_path.display(), e))?;
let response_content_length = response.content_length();
let progress_bar_filename = final_path
.file_name()
.unwrap_or_default()
.to_string_lossy()
.into_owned();
let progress = create_progress_bar(
quiet_mode,
progress_bar_filename,
response_content_length,
false,
);
let mut source = response.take(response_content_length.unwrap_or(u64::MAX));
let mut buffered_reader = progress.wrap_read(&mut source);
let mut buffer = [0u8; 8192];
let mut downloaded: u64 = 0;
let started_at = Instant::now();
loop {
let n = buffered_reader.read(&mut buffer)?;
if n == 0 {
break;
}
dest.write_all(&buffer[..n])?;
downloaded += n as u64;
if let Some(total) = response_content_length {
if let Some(cb) = status_callback {
let percent = downloaded as f64 / total.max(1) as f64 * 100.0;
cb(format!(
"PROGRESS: {:.1}% ({}/{})",
percent, downloaded, total
));
}
}
throttle_download(downloaded, started_at, optimizer.speed_limit);
}
progress.finish_with_message("Download completed\n");
if is_iso && options.verify_iso {
verify_file_sha256(
&final_path,
options.expected_sha256.as_deref(),
status_callback,
)?;
} else if let Some(expected) = options.expected_sha256.as_deref() {
verify_file_sha256(&final_path, Some(expected), status_callback)?;
}
Ok(())
}
pub fn parse_content_disposition_filename(header_value: &str) -> Option<String> {
let mut plain: Option<String> = None;
for part in header_value.split(';') {
let part = part.trim();
if let Some(val) = part.strip_prefix("filename*=") {
let val = val.trim().trim_matches('"');
let mut parts = val.splitn(3, '\'');
let _charset = parts.next();
let _language = parts.next();
if let Some(encoded) = parts.next() {
if let Ok(decoded) = urlencoding::decode(encoded) {
let name = decoded.into_owned();
if !name.is_empty() {
return Some(name);
}
}
}
}
if plain.is_none() {
if let Some(val) = part.strip_prefix("filename=") {
let name = val.trim().trim_matches('"').to_string();
if !name.is_empty() {
plain = Some(name);
}
}
}
}
plain
}
fn throttle_download(downloaded: u64, started_at: Instant, speed_limit: Option<u64>) {
let Some(limit) = speed_limit else { return };
if limit == 0 {
return;
}
let expected_elapsed = Duration::from_secs_f64(downloaded as f64 / limit as f64);
let actual_elapsed = started_at.elapsed();
if expected_elapsed > actual_elapsed {
std::thread::sleep(expected_elapsed - actual_elapsed);
}
}
pub fn verify_iso_integrity(
path: &Path,
callback: Option<&(dyn Fn(String) + Send + Sync)>,
) -> Result<(), Box<dyn Error + Send + Sync>> {
verify_file_sha256(path, None, callback).map(|_| ())
}
pub fn verify_file_sha256(
path: &Path,
expected_hash: Option<&str>,
callback: Option<&(dyn Fn(String) + Send + Sync)>,
) -> Result<String, Box<dyn Error + Send + Sync>> {
let send = |msg: &str| {
if let Some(cb) = callback {
cb(msg.to_string());
}
};
send("Calculating SHA256 hash... (this may take a while for large ISOs)");
let mut file = File::open(path)?;
let mut hasher = sha2::Sha256::new();
let mut buffer = [0; 8192];
loop {
let n = file.read(&mut buffer)?;
if n == 0 {
break;
}
hasher.update(&buffer[..n]);
}
let hash = hex::encode(hasher.finalize());
send("Integrity check finished.");
send(&format!("SHA256: {}", hash));
if let Some(expected_hash) = expected_hash {
let expected_hash = expected_hash.trim().to_ascii_lowercase();
if hash != expected_hash {
return Err(
format!("SHA256 mismatch: expected {}, got {}", expected_hash, hash).into(),
);
}
send("SHA256 matches expected hash.");
}
Ok(hash)
}