1use colored::*;
2use mlua::{self, Lua, Table};
3use std::path::{Path, PathBuf};
4use zoi_core::utils;
5
6use ar::Archive as ArArchive;
7use flate2::read::GzDecoder;
8use sevenz_rust;
9use std::fs;
10use xz2::read::XzDecoder;
11use zip::ZipArchive;
12use zstd::stream::read::Decoder as ZstdDecoder;
13pub fn add_extract_util(lua: &Lua, quiet: bool) -> Result<(), mlua::Error> {
14 let extract_fn =
15 lua.create_function(move |lua, (source, out_name): (String, Option<String>)| {
16 let build_dir_str: String = lua.globals().get("BUILD_DIR")?;
17 let build_dir = Path::new(&build_dir_str);
18
19 let archive_file = if source.starts_with("http") {
20 if source.starts_with("http://") && !quiet {
21 println!("{}: downloading over insecure HTTP: {}", "Warning:".yellow(), source);
22 }
23 if !quiet {
24 println!("Downloading: {}", source);
25 }
26 let file_name = source.split('/').next_back().unwrap_or("download.tmp");
27 let temp_path = build_dir.join(file_name);
28 let client = utils::get_http_client()
29 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
30 let mut attempt = 0u32;
31 let mut response = loop {
32 attempt += 1;
33 match client.get(&source).send() {
34 Ok(resp) => break resp,
35 Err(e) => {
36 if attempt < 3 {
37 if !quiet {
38 eprintln!("Download failed ({}). Retrying...", e);
39 }
40 zoi_core::utils::retry_backoff_sleep(attempt);
41 continue;
42 } else {
43 return Err(mlua::Error::RuntimeError(e.to_string()));
44 }
45 }
46 }
47 };
48
49 if !response.status().is_success() {
50 return Err(mlua::Error::RuntimeError(format!("Failed to download {}: {}", source, response.status())));
51 }
52
53 let mut temp_file = fs::File::create(&temp_path)
54 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
55 std::io::copy(&mut response, &mut temp_file)
56 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
57
58 temp_path
59 } else {
60 PathBuf::from(source)
61 };
62
63 let out_dir_name = out_name.unwrap_or_else(|| "extracted".to_string());
64 let out_dir = build_dir.join(&out_dir_name);
65
66 if !out_dir.starts_with(build_dir) || out_dir == build_dir {
67 return Err(mlua::Error::RuntimeError(format!(
68 "Invalid output directory: {}. Extraction must be into a subdirectory of the build directory.",
69 out_dir_name
70 )));
71 }
72
73 fs::create_dir_all(&out_dir).map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
74
75 if !quiet {
76 println!(
77 "Extracting {} to {}",
78 archive_file.display(),
79 out_dir.display()
80 );
81 }
82
83 let file = fs::File::open(&archive_file)
84 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
85
86 let archive_path_str = archive_file.to_string_lossy();
87
88 if archive_path_str.ends_with(".zip") {
89 let mut archive =
90 ZipArchive::new(file).map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
91 archive
92 .extract(&out_dir)
93 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
94 } else if archive_path_str.ends_with(".tar.gz") || archive_path_str.ends_with(".tgz") {
95 let tar_gz = GzDecoder::new(file);
96 let mut archive = tar::Archive::new(tar_gz);
97 archive
98 .unpack(&out_dir)
99 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
100 } else if archive_path_str.ends_with(".tar.zst") {
101 let tar_zst =
102 ZstdDecoder::new(file).map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
103 let mut archive = tar::Archive::new(tar_zst);
104 archive
105 .unpack(&out_dir)
106 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
107 } else if archive_path_str.ends_with(".tar.xz") {
108 let tar_xz = XzDecoder::new(file);
109 let mut archive = tar::Archive::new(tar_xz);
110 archive
111 .unpack(&out_dir)
112 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
113 } else if archive_path_str.ends_with(".7z") {
114 sevenz_rust::decompress_file(&archive_file, &out_dir)
115 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
116 } else if archive_path_str.ends_with(".dmg") {
117 if !cfg!(target_os = "macos") {
118 return Err(mlua::Error::RuntimeError(
119 "Extracting .dmg files is only supported on macOS.".to_string(),
120 ));
121 }
122 let output = std::process::Command::new("hdiutil")
123 .arg("attach")
124 .arg("-nobrowse")
125 .arg("-readonly")
126 .arg(&archive_file)
127 .output()
128 .map_err(|e| mlua::Error::RuntimeError(format!("Failed to execute hdiutil: {}", e)))?;
129 if !output.status.success() {
130 let stderr = String::from_utf8_lossy(&output.stderr);
131 return Err(mlua::Error::RuntimeError(format!("hdiutil failed: {}", stderr)));
132 }
133 let output_str = String::from_utf8_lossy(&output.stdout);
134 let mut mount_point = None;
135 for line in output_str.lines() {
136 if line.contains("/Volumes/")
137 && let Some(idx) = line.find("/Volumes/") {
138 mount_point = Some(line[idx..].trim().to_string());
139 break;
140 }
141 }
142 let mount_point = mount_point.ok_or_else(|| {
143 mlua::Error::RuntimeError("Failed to parse mount point from hdiutil output.".to_string())
144 })?;
145 let mount_path = std::path::Path::new(&mount_point);
146 if let Err(e) = zoi_core::utils::copy_dir_all(mount_path, &out_dir) {
147 let _ = std::process::Command::new("hdiutil").arg("detach").arg(&mount_point).status();
148 return Err(mlua::Error::RuntimeError(format!("Failed to copy contents from dmg: {}", e)));
149 }
150 let detach_status = std::process::Command::new("hdiutil")
151 .arg("detach")
152 .arg(&mount_point)
153 .status()
154 .map_err(|e| mlua::Error::RuntimeError(format!("Failed to execute hdiutil detach: {}", e)))?;
155 if !detach_status.success() {
156 eprintln!("Warning: failed to detach dmg volume at {}", mount_point);
157 }
158 } else if archive_path_str.ends_with(".pkg") {
159 if !cfg!(target_os = "macos") {
160 return Err(mlua::Error::RuntimeError(
161 "Extracting .pkg files natively is only supported on macOS.".to_string(),
162 ));
163 }
164 let temp_extract_dir = out_dir.join(".pkg_extract_tmp");
165 let status = std::process::Command::new("pkgutil")
166 .arg("--expand-full")
167 .arg(&archive_file)
168 .arg(&temp_extract_dir)
169 .status()
170 .map_err(|e| mlua::Error::RuntimeError(format!("Failed to execute pkgutil: {}", e)))?;
171 if !status.success() {
172 return Err(mlua::Error::RuntimeError("pkgutil failed to expand the package.".to_string()));
173 }
174 zoi_core::utils::copy_dir_all(&temp_extract_dir, &out_dir)
175 .map_err(|e| mlua::Error::RuntimeError(format!("Failed to copy pkg contents: {}", e)))?;
176 let _ = fs::remove_dir_all(&temp_extract_dir);
177
178 } else if archive_path_str.ends_with(".rar") {
179 if zoi_core::utils::command_exists("unrar") {
180 let status = std::process::Command::new("unrar")
181 .arg("x")
182 .arg("-y")
183 .arg(&archive_file)
184 .arg(&out_dir)
185 .status()
186 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
187 if !status.success() {
188 return Err(mlua::Error::RuntimeError("unrar failed".to_string()));
189 }
190 } else {
191 return Err(mlua::Error::RuntimeError(
192 "unrar command not found. Please install unrar to extract .rar files."
193 .to_string(),
194 ));
195 }
196 } else if archive_path_str.ends_with(".deb") {
197 let mut ar = ArArchive::new(file);
198 while let Some(entry_result) = ar.next_entry() {
199 let mut entry =
200 entry_result.map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
201 let name = String::from_utf8_lossy(entry.header().identifier())
202 .trim()
203 .trim_end_matches('/')
204 .to_string();
205 if name.starts_with("data.tar") {
206 let temp_data_path = build_dir.join(&name);
207 let mut temp_file = fs::File::create(&temp_data_path)
208 .map_err(|e| mlua::Error::RuntimeError(format!("Failed to create temp file for {}: {}", name, e)))?;
209 std::io::copy(&mut entry, &mut temp_file)
210 .map_err(|e| mlua::Error::RuntimeError(format!("Failed to copy entry data for {}: {}", name, e)))?;
211
212 let data_file = fs::File::open(&temp_data_path)
213 .map_err(|e| mlua::Error::RuntimeError(format!("Failed to reopen temp file for {}: {}", name, e)))?;
214 if name.ends_with(".gz") {
215 let mut archive = tar::Archive::new(GzDecoder::new(data_file));
216 archive
217 .unpack(&out_dir)
218 .map_err(|e| mlua::Error::RuntimeError(format!("Failed to unpack {}: {}", name, e)))?;
219 } else if name.ends_with(".xz") {
220 let mut archive = tar::Archive::new(XzDecoder::new(data_file));
221 archive
222 .unpack(&out_dir)
223 .map_err(|e| mlua::Error::RuntimeError(format!("Failed to unpack {}: {}", name, e)))?;
224 } else if name.ends_with(".zst") {
225 let mut archive = tar::Archive::new(
226 ZstdDecoder::new(data_file)
227 .map_err(|e| mlua::Error::RuntimeError(format!("Failed to initialize zstd for {}: {}", name, e)))?,
228 );
229 archive
230 .unpack(&out_dir)
231 .map_err(|e| mlua::Error::RuntimeError(format!("Failed to unpack {}: {}", name, e)))?;
232 }
233 fs::remove_file(temp_data_path).ok();
234 }
235 }
236 } else {
237 return Err(mlua::Error::RuntimeError(format!(
238 "Unsupported archive format for file: {}",
239 archive_path_str
240 )));
241 }
242
243 Ok(())
244 })?;
245
246 let utils_table: Table = lua.globals().get("UTILS")?;
247 utils_table.set("EXTRACT", extract_fn)?;
248
249 Ok(())
250}
251
252pub fn add_archive_util(lua: &Lua) -> Result<(), mlua::Error> {
253 let archive_table = lua.create_table()?;
254
255 let list_fn = lua.create_function(|lua, path: String| {
256 let p = Path::new(&path);
257 let actual_path = if p.exists() {
258 p.to_path_buf()
259 } else if let Ok(build_dir) = lua.globals().get::<String>("BUILD_DIR") {
260 Path::new(&build_dir).join(p)
261 } else {
262 p.to_path_buf()
263 };
264
265 let file = fs::File::open(&actual_path).map_err(|e| {
266 mlua::Error::RuntimeError(format!("Failed to open archive {:?}: {}", actual_path, e))
267 })?;
268 let mut files = Vec::new();
269
270 if path.ends_with(".zip") {
271 let mut archive =
272 ZipArchive::new(file).map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
273 for i in 0..archive.len() {
274 let file = archive
275 .by_index(i)
276 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
277 files.push(file.name().to_string());
278 }
279 } else if path.ends_with(".tar.gz") || path.ends_with(".tgz") {
280 let tar_gz = GzDecoder::new(file);
281 let mut archive = tar::Archive::new(tar_gz);
282 for entry in archive
283 .entries()
284 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?
285 {
286 let entry = entry.map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
287 files.push(
288 entry
289 .path()
290 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?
291 .to_string_lossy()
292 .to_string(),
293 );
294 }
295 } else if path.ends_with(".tar.zst") {
296 let tar_zst =
297 ZstdDecoder::new(file).map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
298 let mut archive = tar::Archive::new(tar_zst);
299 for entry in archive
300 .entries()
301 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?
302 {
303 let entry = entry.map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
304 files.push(
305 entry
306 .path()
307 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?
308 .to_string_lossy()
309 .to_string(),
310 );
311 }
312 } else if path.ends_with(".tar.xz") {
313 let tar_xz = XzDecoder::new(file);
314 let mut archive = tar::Archive::new(tar_xz);
315 for entry in archive
316 .entries()
317 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?
318 {
319 let entry = entry.map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
320 files.push(
321 entry
322 .path()
323 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?
324 .to_string_lossy()
325 .to_string(),
326 );
327 }
328 } else if path.ends_with(".7z") {
329 let file =
330 fs::File::open(&path).map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
331 let len = file
332 .metadata()
333 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?
334 .len();
335 let reader = sevenz_rust::SevenZReader::new(file, len, sevenz_rust::Password::empty())
336 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
337 for entry in &reader.archive().files {
338 files.push(entry.name.to_string());
339 }
340 } else if path.ends_with(".rar") {
341 if zoi_core::utils::command_exists("unrar") {
342 let output = std::process::Command::new("unrar")
343 .arg("lb")
344 .arg(&path)
345 .output()
346 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
347 if output.status.success() {
348 let list = String::from_utf8_lossy(&output.stdout);
349 for line in list.lines() {
350 files.push(line.to_string());
351 }
352 }
353 }
354 } else if path.ends_with(".deb") {
355 let mut ar = ArArchive::new(file);
356 while let Some(entry_result) = ar.next_entry() {
357 let entry = entry_result.map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
358 let header = entry.header();
359 files.push(String::from_utf8_lossy(header.identifier()).to_string());
360 }
361 } else {
362 return Err(mlua::Error::RuntimeError(format!(
363 "Unsupported archive format: {}",
364 path
365 )));
366 }
367
368 Ok(files)
369 })?;
370 archive_table.set("list", list_fn)?;
371
372 let utils_table: Table = lua.globals().get("UTILS")?;
373 utils_table.set("ARCHIVE", archive_table)?;
374
375 Ok(())
376}