kget/
download.rs

1use reqwest::blocking::Client;
2use reqwest::header::{CONTENT_LENGTH, CONTENT_TYPE};
3use std::fs::File;
4use std::time::Duration;
5use std::error::Error;
6use crate::progress::create_progress_bar;
7use crate::utils::{self, print};
8use humansize::{format_size, DECIMAL};
9use mime::Mime;
10use std::io::{Read, Write};
11use std::path::{Path, PathBuf};
12use sha2::Digest;
13use crate::config::ProxyConfig;
14use crate::optimization::Optimizer;
15use crate::DownloadOptions;
16
17const MAX_RETRIES: u32 = 3;
18const RETRY_DELAY: Duration = Duration::from_secs(2);
19
20pub fn check_disk_space(path: &Path, required_size: u64) -> Result<(), Box<dyn Error + Send + Sync>> {
21    let dir = path.parent().unwrap_or(Path::new("."));
22    let available_space = fs2::available_space(dir)?;
23    
24    if available_space < required_size {
25        return Err(format!(
26            "Insufficient disk space. Required: {}, Available: {}", 
27            format_size(required_size, DECIMAL),
28            format_size(available_space, DECIMAL)
29        ).into());
30    }
31    Ok(())
32}
33
34pub fn validate_filename(filename: &str) -> Result<(), Box<dyn Error + Send + Sync>> {
35    if filename.contains(std::path::MAIN_SEPARATOR) {
36        return Err("Filename cannot contain directory separators".into());
37    }
38    if filename.is_empty() {
39        return Err("Filename cannot be empty".into());
40    }
41    Ok(())
42}
43
44pub fn download(
45    target: &str,
46    proxy: ProxyConfig,
47    optimizer: Optimizer,
48    options: DownloadOptions,
49    status_callback: Option<&(dyn Fn(String) + Send + Sync)>,
50) -> Result<(), Box<dyn Error + Send + Sync>> {
51    let quiet_mode = options.quiet_mode;
52   
53
54    let mut client_builder = Client::builder()
55        .timeout(Duration::from_secs(30))
56        .no_gzip()
57        .no_deflate();
58
59    if proxy.enabled {
60        if let Some(proxy_url) = &proxy.url {
61            let proxy_client = match proxy.proxy_type {
62                crate::config::ProxyType::Http => reqwest::Proxy::http(proxy_url),
63                crate::config::ProxyType::Https => reqwest::Proxy::https(proxy_url),
64                crate::config::ProxyType::Socks5 => reqwest::Proxy::all(proxy_url),
65            };
66            if let Ok(mut proxy_client) = proxy_client {
67                if let (Some(username), Some(password)) = (&proxy.username, &proxy.password) {
68                    proxy_client = proxy_client.basic_auth(username, password);
69                }
70                client_builder = client_builder.proxy(proxy_client);
71            }
72        }
73    }
74
75    let client = client_builder.build()?;
76
77    let mut retries = 0;
78    let response = loop {
79        match client.get(target).send() {
80            Ok(resp) => break resp,
81            Err(e) => {
82                retries += 1;
83                if retries >= MAX_RETRIES {
84                    return Err(format!("Failed after {} attempts: {}", MAX_RETRIES, e).into());
85                }
86                print(&format!("Attempt {} failed, retrying in {} seconds...", 
87                    retries, RETRY_DELAY.as_secs()), quiet_mode);
88                std::thread::sleep(RETRY_DELAY);
89            }
90        }
91    };
92    
93    print(
94        &format!("HTTP request sent... {}", response.status()),
95        quiet_mode
96    );
97
98    if !response.status().is_success() {
99        return Err(format!("HTTP error: {}", response.status()).into());
100    }
101
102    let content_length = response.headers()
103        .get(CONTENT_LENGTH)
104        .and_then(|ct_len| ct_len.to_str().ok())
105        .and_then(|s| s.parse::<u64>().ok());
106
107    let content_type = response.headers()
108        .get(CONTENT_TYPE)
109        .and_then(|ct| ct.to_str().ok())
110        .and_then(|s| s.parse::<Mime>().ok());
111
112    if let Some(len) = content_length {
113        print(
114            &format!("Length: {} ({})", 
115                len, 
116                format_size(len, DECIMAL)
117            ), 
118            quiet_mode
119        );
120    } else {
121        print("Length: unknown", quiet_mode);
122    }
123
124    if let Some(ref ct) = content_type {
125        print(&format!("Type: {}", ct), quiet_mode);
126    }
127
128    let is_iso = target.to_lowercase().ends_with(".iso") 
129        || content_type.as_ref().map_or(false, |ct| ct.essence_str() == "application/x-iso9001" || ct.essence_str() == "application/x-cd-image");
130
131    if is_iso {
132        print("ISO file detected. Ensuring raw download to prevent corruption...", quiet_mode);
133    }
134
135    let mut tentative_path: PathBuf;
136
137    if let Some(output_arg_str) = options.output_path { 
138        let user_path = PathBuf::from(output_arg_str.clone());
139
140        let is_target_dir = user_path.is_dir() || 
141                              output_arg_str.ends_with(std::path::MAIN_SEPARATOR);
142
143        if is_target_dir {
144            let base_filename = utils::get_filename_from_url_or_default(target, "downloaded_file");
145            validate_filename(&base_filename)?;
146            tentative_path = user_path.join(base_filename);
147        } else {
148            if let Some(file_name_osstr) = user_path.file_name() {
149                if let Some(file_name_str) = file_name_osstr.to_str() {
150                    if file_name_str.is_empty() {
151                        return Err(format!("Invalid output path, does not specify a file name: {}", user_path.display()).into());
152                    }
153                    validate_filename(file_name_str)?;
154                } else {
155                    return Err("Output filename contains invalid characters (non-UTF-8)".into());
156                }
157            } else {
158                return Err(format!("Invalid output path, does not specify a file name: {}", user_path.display()).into());
159            }
160            tentative_path = user_path;
161        }
162    } else {
163        let base_filename = utils::get_filename_from_url_or_default(target, "downloaded_file");
164        validate_filename(&base_filename)?;
165        tentative_path = PathBuf::from(base_filename);
166    }
167
168    let final_path: PathBuf = if tentative_path.is_absolute() {
169        tentative_path
170    } else {
171        let current_dir = std::env::current_dir()
172            .map_err(|e| format!("Failed to get current directory: {}", e))?;
173        current_dir.join(tentative_path)
174    };
175
176    if let Some(parent_dir) = final_path.parent() {
177        if !parent_dir.as_os_str().is_empty() && parent_dir != Path::new("/") && !parent_dir.exists() {
178            std::fs::create_dir_all(parent_dir)
179                .map_err(|e| format!("Failed to create directory {}: {}", parent_dir.display(), e))?;
180            if !quiet_mode {
181                print(&format!("Created directory: {}", parent_dir.display()), quiet_mode);
182            }
183        }
184    }
185    
186    if !quiet_mode {
187        print(&format!("Saving to: {}", final_path.display()), quiet_mode);
188    }
189
190    if let Some(len) = content_length {
191        check_disk_space(&final_path, len)?;
192    }
193
194    let mut dest = File::create(&final_path).map_err(|e| {
195        format!("Failed to create file {}: {}", final_path.display(), e)
196    })?;
197    
198    let response_content_length = response.content_length();
199    let progress_bar_filename = final_path.file_name().unwrap_or_default().to_string_lossy().into_owned();
200    let progress = create_progress_bar(quiet_mode, progress_bar_filename, response_content_length, false);
201    
202    let mut source = response.take(response_content_length.unwrap_or(u64::MAX));
203    let mut buffered_reader = progress.wrap_read(&mut source);
204    
205    // Stream data instead of reading all into memory
206    let mut buffer = [0u8; 8192];
207    loop {
208        let n = buffered_reader.read(&mut buffer)?;
209        if n == 0 { break; }
210        dest.write_all(&buffer[..n])?;
211    }
212    
213    progress.finish_with_message("Download completed\n");
214
215    
216    if is_iso && options.verify_iso {
217        verify_iso_integrity(&final_path, status_callback)?;
218    }
219
220    Ok(())
221}
222
223
224pub fn verify_iso_integrity(path: &Path, callback: Option<&(dyn Fn(String) + Send + Sync)>) -> Result<(), Box<dyn Error + Send + Sync>> {
225    let msg = "Calculating SHA256 hash... (this may take a while for large ISOs)";
226    if let Some(cb) = callback { cb(msg.to_string()); }
227    println!("{}", msg);
228
229    let mut file = File::open(path)?;
230    let mut hasher = sha2::Sha256::new();
231    let mut buffer = [0; 8192];
232    loop {
233        let n = file.read(&mut buffer)?;
234        if n == 0 { break; }
235        hasher.update(&buffer[..n]);
236    }
237    let hash = hex::encode(hasher.finalize());
238    
239    let msg_done = "Integrity check finished.";
240    if let Some(cb) = callback { cb(msg_done.to_string()); }
241    println!("{}", msg_done);
242
243    let msg_hash = format!("SHA256: {}", hash);
244    if let Some(cb) = callback { cb(msg_hash); }
245    println!("SHA256: {}", hash);
246    
247    println!("You can compare this hash with the one provided by the source.");
248    Ok(())
249}