1use crate::DownloadOptions;
31use crate::config::ProxyConfig;
32use crate::optimization::Optimizer;
33use crate::progress::create_progress_bar;
34use crate::utils::{self, print};
35use humansize::{DECIMAL, format_size};
36use mime::Mime;
37use reqwest::blocking::Client;
38use reqwest::header::{CONTENT_DISPOSITION, CONTENT_LENGTH, CONTENT_TYPE};
39use sha2::Digest;
40use std::error::Error;
41use std::fs::File;
42use std::io::{Read, Write};
43use std::path::{Path, PathBuf};
44use std::time::{Duration, Instant};
45
46const MAX_RETRIES: u32 = 3;
47const RETRY_DELAY: Duration = Duration::from_secs(2);
48
49pub fn check_disk_space(
60 path: &Path,
61 required_size: u64,
62) -> Result<(), Box<dyn Error + Send + Sync>> {
63 let dir = path.parent().unwrap_or(Path::new("."));
64 let available_space = fs2::available_space(dir)?;
65
66 if available_space < required_size {
67 return Err(format!(
68 "Insufficient disk space. Required: {}, Available: {}",
69 format_size(required_size, DECIMAL),
70 format_size(available_space, DECIMAL)
71 )
72 .into());
73 }
74 Ok(())
75}
76
77pub fn validate_filename(filename: &str) -> Result<(), Box<dyn Error + Send + Sync>> {
85 if filename.contains('/') || filename.contains('\\') {
86 return Err("Filename cannot contain directory separators".into());
87 }
88 if filename.is_empty() {
89 return Err("Filename cannot be empty".into());
90 }
91 Ok(())
92}
93
94pub fn download(
134 target: &str,
135 proxy: ProxyConfig,
136 optimizer: Optimizer,
137 options: DownloadOptions,
138 status_callback: Option<&(dyn Fn(String) + Send + Sync)>,
139) -> Result<(), Box<dyn Error + Send + Sync>> {
140 let quiet_mode = options.quiet_mode;
141
142 let mut client_builder = Client::builder()
143 .timeout(Duration::from_secs(30))
144 .user_agent(concat!("KGet/", env!("CARGO_PKG_VERSION")))
145 .no_gzip()
146 .no_deflate();
147
148 if proxy.enabled {
149 if let Some(proxy_url) = &proxy.url {
150 let proxy_client = match proxy.proxy_type {
151 crate::config::ProxyType::Http => reqwest::Proxy::http(proxy_url),
152 crate::config::ProxyType::Https => reqwest::Proxy::https(proxy_url),
153 crate::config::ProxyType::Socks5 => reqwest::Proxy::all(proxy_url),
154 };
155 if let Ok(mut proxy_client) = proxy_client {
156 if let (Some(username), Some(password)) = (&proxy.username, &proxy.password) {
157 proxy_client = proxy_client.basic_auth(username, password);
158 }
159 client_builder = client_builder.proxy(proxy_client);
160 }
161 }
162 }
163
164 let client = client_builder.build()?;
165
166 let mut retries = 0;
167 let response = loop {
168 match client.get(target).send() {
169 Ok(resp) => break resp,
170 Err(e) => {
171 retries += 1;
172 if retries >= MAX_RETRIES {
173 return Err(format!("Failed after {} attempts: {}", MAX_RETRIES, e).into());
174 }
175 print(
176 &format!(
177 "Attempt {} failed, retrying in {} seconds...",
178 retries,
179 RETRY_DELAY.as_secs()
180 ),
181 quiet_mode,
182 );
183 std::thread::sleep(RETRY_DELAY);
184 }
185 }
186 };
187
188 print(
189 &format!("HTTP request sent... {}", response.status()),
190 quiet_mode,
191 );
192
193 if !response.status().is_success() {
194 return Err(format!("HTTP error: {}", response.status()).into());
195 }
196
197 let content_length = response
198 .headers()
199 .get(CONTENT_LENGTH)
200 .and_then(|ct_len| ct_len.to_str().ok())
201 .and_then(|s| s.parse::<u64>().ok());
202
203 let content_type = response
204 .headers()
205 .get(CONTENT_TYPE)
206 .and_then(|ct| ct.to_str().ok())
207 .and_then(|s| s.parse::<Mime>().ok());
208
209 let server_filename = response
210 .headers()
211 .get(CONTENT_DISPOSITION)
212 .and_then(|v| v.to_str().ok())
213 .and_then(parse_content_disposition_filename);
214
215 if let Some(len) = content_length {
216 print(
217 &format!("Length: {} ({})", len, format_size(len, DECIMAL)),
218 quiet_mode,
219 );
220 } else {
221 print("Length: unknown", quiet_mode);
222 }
223
224 if let Some(ref ct) = content_type {
225 print(&format!("Type: {}", ct), quiet_mode);
226 }
227
228 let is_iso = target.to_lowercase().ends_with(".iso")
229 || content_type.as_ref().map_or(false, |ct| {
230 ct.essence_str() == "application/x-iso9001"
231 || ct.essence_str() == "application/x-cd-image"
232 });
233
234 if is_iso {
235 print(
236 "ISO file detected. Ensuring raw download to prevent corruption...",
237 quiet_mode,
238 );
239 }
240
241 let tentative_path: PathBuf;
242
243 if let Some(output_arg_str) = options.output_path {
244 let user_path = PathBuf::from(output_arg_str.clone());
245
246 let is_target_dir =
247 user_path.is_dir() || output_arg_str.ends_with(std::path::MAIN_SEPARATOR);
248
249 if is_target_dir {
250 let base_filename = utils::get_filename_from_url_or_default(target, "downloaded_file");
251 validate_filename(&base_filename)?;
252 tentative_path = user_path.join(base_filename);
253 } else {
254 if let Some(file_name_osstr) = user_path.file_name() {
255 if let Some(file_name_str) = file_name_osstr.to_str() {
256 if file_name_str.is_empty() {
257 return Err(format!(
258 "Invalid output path, does not specify a file name: {}",
259 user_path.display()
260 )
261 .into());
262 }
263 validate_filename(file_name_str)?;
264 } else {
265 return Err("Output filename contains invalid characters (non-UTF-8)".into());
266 }
267 } else {
268 return Err(format!(
269 "Invalid output path, does not specify a file name: {}",
270 user_path.display()
271 )
272 .into());
273 }
274 tentative_path = user_path;
275 }
276 } else {
277 let base_filename = if let Some(ref name) = server_filename {
278 name.clone()
279 } else {
280 utils::get_filename_from_url_or_default(target, "downloaded_file")
281 };
282 validate_filename(&base_filename)?;
283 tentative_path = PathBuf::from(base_filename);
284 }
285
286 let final_path: PathBuf = if tentative_path.is_absolute() {
287 tentative_path
288 } else {
289 let current_dir = std::env::current_dir()
290 .map_err(|e| format!("Failed to get current directory: {}", e))?;
291 current_dir.join(tentative_path)
292 };
293
294 if let Some(parent_dir) = final_path.parent() {
295 if !parent_dir.as_os_str().is_empty()
296 && parent_dir != Path::new("/")
297 && !parent_dir.exists()
298 {
299 std::fs::create_dir_all(parent_dir).map_err(|e| {
300 format!("Failed to create directory {}: {}", parent_dir.display(), e)
301 })?;
302 if !quiet_mode {
303 print(
304 &format!("Created directory: {}", parent_dir.display()),
305 quiet_mode,
306 );
307 }
308 }
309 }
310
311 if !quiet_mode {
312 print(&format!("Saving to: {}", final_path.display()), quiet_mode);
313 }
314
315 if let Some(len) = content_length {
316 check_disk_space(&final_path, len)?;
317 }
318
319 let mut dest = File::create(&final_path)
320 .map_err(|e| format!("Failed to create file {}: {}", final_path.display(), e))?;
321
322 let response_content_length = response.content_length();
323 let progress_bar_filename = final_path
324 .file_name()
325 .unwrap_or_default()
326 .to_string_lossy()
327 .into_owned();
328 let progress = create_progress_bar(
329 quiet_mode,
330 progress_bar_filename,
331 response_content_length,
332 false,
333 );
334
335 let mut source = response.take(response_content_length.unwrap_or(u64::MAX));
336 let mut buffered_reader = progress.wrap_read(&mut source);
337
338 let mut buffer = [0u8; 8192];
340 let mut downloaded: u64 = 0;
341 let started_at = Instant::now();
342 loop {
343 let n = buffered_reader.read(&mut buffer)?;
344 if n == 0 {
345 break;
346 }
347 dest.write_all(&buffer[..n])?;
348 downloaded += n as u64;
349
350 if let Some(total) = response_content_length {
351 if let Some(cb) = status_callback {
352 let percent = downloaded as f64 / total.max(1) as f64 * 100.0;
353 cb(format!("PROGRESS: {:.1}% ({}/{})", percent, downloaded, total));
354 }
355 }
356
357 throttle_download(downloaded, started_at, optimizer.speed_limit);
358 }
359
360 progress.finish_with_message("Download completed\n");
361
362 if is_iso && options.verify_iso {
363 verify_file_sha256(
364 &final_path,
365 options.expected_sha256.as_deref(),
366 status_callback,
367 )?;
368 } else if let Some(expected) = options.expected_sha256.as_deref() {
369 verify_file_sha256(&final_path, Some(expected), status_callback)?;
370 }
371
372 Ok(())
373}
374
375pub fn parse_content_disposition_filename(header_value: &str) -> Option<String> {
380 let mut plain: Option<String> = None;
381
382 for part in header_value.split(';') {
383 let part = part.trim();
384
385 if let Some(val) = part.strip_prefix("filename*=") {
387 let val = val.trim().trim_matches('"');
388 let mut parts = val.splitn(3, '\'');
390 let _charset = parts.next();
391 let _language = parts.next();
392 if let Some(encoded) = parts.next() {
393 if let Ok(decoded) = urlencoding::decode(encoded) {
394 let name = decoded.into_owned();
395 if !name.is_empty() {
396 return Some(name);
397 }
398 }
399 }
400 }
401
402 if plain.is_none() {
404 if let Some(val) = part.strip_prefix("filename=") {
405 let name = val.trim().trim_matches('"').to_string();
406 if !name.is_empty() {
407 plain = Some(name);
408 }
409 }
410 }
411 }
412
413 plain
414}
415
416fn throttle_download(downloaded: u64, started_at: Instant, speed_limit: Option<u64>) {
417 let Some(limit) = speed_limit else { return };
418 if limit == 0 {
419 return;
420 }
421
422 let expected_elapsed = Duration::from_secs_f64(downloaded as f64 / limit as f64);
423 let actual_elapsed = started_at.elapsed();
424 if expected_elapsed > actual_elapsed {
425 std::thread::sleep(expected_elapsed - actual_elapsed);
426 }
427}
428
429pub fn verify_iso_integrity(
455 path: &Path,
456 callback: Option<&(dyn Fn(String) + Send + Sync)>,
457) -> Result<(), Box<dyn Error + Send + Sync>> {
458 verify_file_sha256(path, None, callback).map(|_| ())
459}
460
461pub fn verify_file_sha256(
463 path: &Path,
464 expected_hash: Option<&str>,
465 callback: Option<&(dyn Fn(String) + Send + Sync)>,
466) -> Result<String, Box<dyn Error + Send + Sync>> {
467 let msg = "Calculating SHA256 hash... (this may take a while for large ISOs)";
468 if let Some(cb) = callback {
469 cb(msg.to_string());
470 }
471 println!("{}", msg);
472
473 let mut file = File::open(path)?;
474 let mut hasher = sha2::Sha256::new();
475 let mut buffer = [0; 8192];
476 loop {
477 let n = file.read(&mut buffer)?;
478 if n == 0 {
479 break;
480 }
481 hasher.update(&buffer[..n]);
482 }
483 let hash = hex::encode(hasher.finalize());
484
485 let msg_done = "Integrity check finished.";
486 if let Some(cb) = callback {
487 cb(msg_done.to_string());
488 }
489 println!("{}", msg_done);
490
491 let msg_hash = format!("SHA256: {}", hash);
492 if let Some(cb) = callback {
493 cb(msg_hash.clone());
494 }
495 println!("SHA256: {}", hash);
496
497 if let Some(expected_hash) = expected_hash {
498 let expected_hash = expected_hash.trim().to_ascii_lowercase();
499 if hash != expected_hash {
500 return Err(
501 format!("SHA256 mismatch: expected {}, got {}", expected_hash, hash).into(),
502 );
503 }
504
505 let msg_match = "SHA256 matches expected hash.";
506 if let Some(cb) = callback {
507 cb(msg_match.to_string());
508 }
509 println!("{}", msg_match);
510 } else {
511 println!("You can compare this hash with the one provided by the source.");
512 }
513
514 Ok(hash)
515}