1use mlua::{self, Lua, Table};
2use std::path::{Path, PathBuf};
3
4use ar::Archive as ArArchive;
5use flate2::read::GzDecoder;
6use sevenz_rust;
7use std::fs;
8use xz2::read::XzDecoder;
9use zip::ZipArchive;
10use zstd::stream::read::Decoder as ZstdDecoder;
11
12pub fn add_extract_util(lua: &Lua, quiet: bool) -> Result<(), mlua::Error> {
22 let extract_fn =
23 lua.create_function(move |lua, (source, out_name): (String, Option<String>)| {
24 let build_dir_str: String = lua.globals().get("BUILD_DIR")?;
25 let build_dir = Path::new(&build_dir_str);
26
27 let archive_file = if source.starts_with("http") {
28 let file_name = source.split('/').next_back().unwrap_or("download.tmp");
29 let temp_path = build_dir.join(file_name);
30 super::download::download_with_progress(&source, &temp_path, quiet)?;
31
32 temp_path
33 } else {
34 PathBuf::from(source)
35 };
36
37 let out_dir_name = out_name.unwrap_or_else(|| "extracted".to_string());
38 let out_dir = build_dir.join(&out_dir_name);
39
40 if !out_dir.starts_with(build_dir) || out_dir == build_dir {
41 return Err(mlua::Error::RuntimeError(format!(
42 "Invalid output directory: {}. Extraction must be into a subdirectory of the build directory.",
43 out_dir_name
44 )));
45 }
46
47 fs::create_dir_all(&out_dir).map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
48
49 if !quiet {
50 println!(
51 "Extracting {} to {}",
52 archive_file.display(),
53 out_dir.display()
54 );
55 }
56
57 let file = fs::File::open(&archive_file)
58 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
59
60 let archive_path_str = archive_file.to_string_lossy();
61
62 if archive_path_str.ends_with(".zip") {
63 let mut archive =
64 ZipArchive::new(file).map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
65 archive
66 .extract(&out_dir)
67 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
68 } else if archive_path_str.ends_with(".tar.gz") || archive_path_str.ends_with(".tgz") {
69 let tar_gz = GzDecoder::new(file);
70 let mut archive = tar::Archive::new(tar_gz);
71 archive
72 .unpack(&out_dir)
73 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
74 } else if archive_path_str.ends_with(".tar.zst")
75 || archive_path_str.ends_with(".zpa")
76 || archive_path_str.ends_with(".zsa")
77 {
78 let tar_zst =
79 ZstdDecoder::new(file).map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
80 let mut archive = tar::Archive::new(tar_zst);
81 archive
82 .unpack(&out_dir)
83 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
84 } else if archive_path_str.ends_with(".tar.xz") {
85 let tar_xz = XzDecoder::new(file);
86 let mut archive = tar::Archive::new(tar_xz);
87 archive
88 .unpack(&out_dir)
89 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
90 } else if archive_path_str.ends_with(".7z") {
91 sevenz_rust::decompress_file(&archive_file, &out_dir)
92 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
93 } else if archive_path_str.ends_with(".dmg") {
94 if !cfg!(target_os = "macos") {
95 return Err(mlua::Error::RuntimeError(
96 "Extracting .dmg files is only supported on macOS.".to_string(),
97 ));
98 }
99 let output = std::process::Command::new("hdiutil")
100 .arg("attach")
101 .arg("-nobrowse")
102 .arg("-readonly")
103 .arg(&archive_file)
104 .output()
105 .map_err(|e| mlua::Error::RuntimeError(format!("Failed to execute hdiutil: {}", e)))?;
106 if !output.status.success() {
107 let stderr = String::from_utf8_lossy(&output.stderr);
108 return Err(mlua::Error::RuntimeError(format!("hdiutil failed: {}", stderr)));
109 }
110 let output_str = String::from_utf8_lossy(&output.stdout);
111 let mut mount_point = None;
112 for line in output_str.lines() {
113 if line.contains("/Volumes/")
114 && let Some(idx) = line.find("/Volumes/") {
115 mount_point = Some(line[idx..].trim().to_string());
116 break;
117 }
118 }
119 let mount_point = mount_point.ok_or_else(|| {
120 mlua::Error::RuntimeError("Failed to parse mount point from hdiutil output.".to_string())
121 })?;
122 let mount_path = std::path::Path::new(&mount_point);
123 if let Err(e) = zoi_core::utils::copy_dir_all(mount_path, &out_dir) {
124 let _ = std::process::Command::new("hdiutil").arg("detach").arg(&mount_point).status();
125 return Err(mlua::Error::RuntimeError(format!("Failed to copy contents from dmg: {}", e)));
126 }
127 let detach_status = std::process::Command::new("hdiutil")
128 .arg("detach")
129 .arg(&mount_point)
130 .status()
131 .map_err(|e| mlua::Error::RuntimeError(format!("Failed to execute hdiutil detach: {}", e)))?;
132 if !detach_status.success() {
133 eprintln!("Warning: failed to detach dmg volume at {}", mount_point);
134 }
135 } else if archive_path_str.ends_with(".pkg") {
136 if !cfg!(target_os = "macos") {
137 return Err(mlua::Error::RuntimeError(
138 "Extracting .pkg files natively is only supported on macOS.".to_string(),
139 ));
140 }
141 let temp_extract_dir = out_dir.join(".pkg_extract_tmp");
142 let status = std::process::Command::new("pkgutil")
143 .arg("--expand-full")
144 .arg(&archive_file)
145 .arg(&temp_extract_dir)
146 .status()
147 .map_err(|e| mlua::Error::RuntimeError(format!("Failed to execute pkgutil: {}", e)))?;
148 if !status.success() {
149 return Err(mlua::Error::RuntimeError("pkgutil failed to expand the package.".to_string()));
150 }
151 zoi_core::utils::copy_dir_all(&temp_extract_dir, &out_dir)
152 .map_err(|e| mlua::Error::RuntimeError(format!("Failed to copy pkg contents: {}", e)))?;
153 let _ = fs::remove_dir_all(&temp_extract_dir);
154
155 } else if archive_path_str.ends_with(".rar") {
156 if zoi_core::utils::command_exists("unrar") {
157 let status = std::process::Command::new("unrar")
158 .arg("x")
159 .arg("-y")
160 .arg(&archive_file)
161 .arg(&out_dir)
162 .status()
163 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
164 if !status.success() {
165 return Err(mlua::Error::RuntimeError("unrar failed".to_string()));
166 }
167 } else {
168 return Err(mlua::Error::RuntimeError(
169 "unrar command not found. Please install unrar to extract .rar files."
170 .to_string(),
171 ));
172 }
173 } else if archive_path_str.ends_with(".deb") {
174 let mut ar = ArArchive::new(file);
175 while let Some(entry_result) = ar.next_entry() {
176 let mut entry =
177 entry_result.map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
178 let name = String::from_utf8_lossy(entry.header().identifier())
179 .trim()
180 .trim_end_matches('/')
181 .to_string();
182 if name.starts_with("data.tar") {
183 let temp_data_path = build_dir.join(&name);
184 let mut temp_file = fs::File::create(&temp_data_path)
185 .map_err(|e| mlua::Error::RuntimeError(format!("Failed to create temp file for {}: {}", name, e)))?;
186 std::io::copy(&mut entry, &mut temp_file)
187 .map_err(|e| mlua::Error::RuntimeError(format!("Failed to copy entry data for {}: {}", name, e)))?;
188
189 let data_file = fs::File::open(&temp_data_path)
190 .map_err(|e| mlua::Error::RuntimeError(format!("Failed to reopen temp file for {}: {}", name, e)))?;
191 if name.ends_with(".gz") {
192 let mut archive = tar::Archive::new(GzDecoder::new(data_file));
193 archive
194 .unpack(&out_dir)
195 .map_err(|e| mlua::Error::RuntimeError(format!("Failed to unpack {}: {}", name, e)))?;
196 } else if name.ends_with(".xz") {
197 let mut archive = tar::Archive::new(XzDecoder::new(data_file));
198 archive
199 .unpack(&out_dir)
200 .map_err(|e| mlua::Error::RuntimeError(format!("Failed to unpack {}: {}", name, e)))?;
201 } else if name.ends_with(".zst") {
202 let mut archive = tar::Archive::new(
203 ZstdDecoder::new(data_file)
204 .map_err(|e| mlua::Error::RuntimeError(format!("Failed to initialize zstd for {}: {}", name, e)))?,
205 );
206 archive
207 .unpack(&out_dir)
208 .map_err(|e| mlua::Error::RuntimeError(format!("Failed to unpack {}: {}", name, e)))?;
209 }
210 fs::remove_file(temp_data_path).ok();
211 }
212 }
213 } else {
214 return Err(mlua::Error::RuntimeError(format!(
215 "Unsupported archive format for file: {}",
216 archive_path_str
217 )));
218 }
219
220 Ok(())
221 })?;
222
223 let utils_table: Table = lua.globals().get("UTILS")?;
224 utils_table.set("EXTRACT", extract_fn)?;
225
226 Ok(())
227}
228
229pub fn add_archive_util(lua: &Lua) -> Result<(), mlua::Error> {
230 let archive_table = lua.create_table()?;
231
232 let list_fn = lua.create_function(|lua, path: String| {
233 let p = Path::new(&path);
234 let actual_path = if p.exists() {
235 p.to_path_buf()
236 } else if let Ok(build_dir) = lua.globals().get::<String>("BUILD_DIR") {
237 Path::new(&build_dir).join(p)
238 } else {
239 p.to_path_buf()
240 };
241
242 let file = fs::File::open(&actual_path).map_err(|e| {
243 mlua::Error::RuntimeError(format!("Failed to open archive {:?}: {}", actual_path, e))
244 })?;
245 let mut files = Vec::new();
246
247 if path.ends_with(".zip") {
248 let mut archive =
249 ZipArchive::new(file).map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
250 for i in 0..archive.len() {
251 let file = archive
252 .by_index(i)
253 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
254 files.push(file.name().to_string());
255 }
256 } else if path.ends_with(".tar.gz") || path.ends_with(".tgz") {
257 let tar_gz = GzDecoder::new(file);
258 let mut archive = tar::Archive::new(tar_gz);
259 for entry in archive
260 .entries()
261 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?
262 {
263 let entry = entry.map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
264 files.push(
265 entry
266 .path()
267 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?
268 .to_string_lossy()
269 .to_string(),
270 );
271 }
272 } else if path.ends_with(".tar.zst") || path.ends_with(".zpa") || path.ends_with(".zsa") {
273 let tar_zst =
274 ZstdDecoder::new(file).map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
275 let mut archive = tar::Archive::new(tar_zst);
276 for entry in archive
277 .entries()
278 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?
279 {
280 let entry = entry.map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
281 files.push(
282 entry
283 .path()
284 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?
285 .to_string_lossy()
286 .to_string(),
287 );
288 }
289 } else if path.ends_with(".tar.xz") {
290 let tar_xz = XzDecoder::new(file);
291 let mut archive = tar::Archive::new(tar_xz);
292 for entry in archive
293 .entries()
294 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?
295 {
296 let entry = entry.map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
297 files.push(
298 entry
299 .path()
300 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?
301 .to_string_lossy()
302 .to_string(),
303 );
304 }
305 } else if path.ends_with(".7z") {
306 let file =
307 fs::File::open(&path).map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
308 let len = file
309 .metadata()
310 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?
311 .len();
312 let reader = sevenz_rust::SevenZReader::new(file, len, sevenz_rust::Password::empty())
313 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
314 for entry in &reader.archive().files {
315 files.push(entry.name.to_string());
316 }
317 } else if path.ends_with(".rar") {
318 if zoi_core::utils::command_exists("unrar") {
319 let output = std::process::Command::new("unrar")
320 .arg("lb")
321 .arg(&path)
322 .output()
323 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
324 if output.status.success() {
325 let list = String::from_utf8_lossy(&output.stdout);
326 for line in list.lines() {
327 files.push(line.to_string());
328 }
329 }
330 }
331 } else if path.ends_with(".deb") {
332 let mut ar = ArArchive::new(file);
333 while let Some(entry_result) = ar.next_entry() {
334 let entry = entry_result.map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
335 let header = entry.header();
336 files.push(String::from_utf8_lossy(header.identifier()).to_string());
337 }
338 } else {
339 return Err(mlua::Error::RuntimeError(format!(
340 "Unsupported archive format: {}",
341 path
342 )));
343 }
344
345 Ok(files)
346 })?;
347 archive_table.set("list", list_fn)?;
348
349 let make_archive_fn = lua.create_function(
350 move |lua, (source, output, algorithm): (mlua::Value, String, Option<String>)| {
351 let algo = algorithm
352 .unwrap_or_else(|| "zst".to_string())
353 .to_lowercase();
354 let build_dir_str: String = lua.globals().get("BUILD_DIR")?;
355 let build_dir = Path::new(&build_dir_str);
356
357 let output_path = if Path::new(&output).is_absolute() {
358 PathBuf::from(&output)
359 } else {
360 build_dir.join(&output)
361 };
362
363 if let Some(parent) = output_path.parent() {
364 fs::create_dir_all(parent).map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
365 }
366
367 let mut source_paths = Vec::new();
368 match source {
369 mlua::Value::String(s) => {
370 let s_str = s.to_str()?;
371 let p = build_dir.join(&*s_str);
372 if p.exists() {
373 source_paths.push((p, s_str.to_string()));
374 } else if Path::new(&*s_str).exists() {
375 source_paths.push((PathBuf::from(s_str.to_string()), s_str.to_string()));
376 } else {
377 return Err(mlua::Error::RuntimeError(format!(
378 "MAKE_ARCHIVE: source path does not exist: {}",
379 &*s_str
380 )));
381 }
382 }
383 mlua::Value::Table(t) => {
384 for val in t.sequence_values::<String>() {
385 let s_str = val?;
386 let p = build_dir.join(&s_str);
387 if p.exists() {
388 source_paths.push((p, s_str.to_string()));
389 } else if Path::new(&s_str).exists() {
390 source_paths.push((PathBuf::from(&s_str), s_str.to_string()));
391 } else {
392 return Err(mlua::Error::RuntimeError(format!(
393 "MAKE_ARCHIVE: source path does not exist: {}",
394 s_str
395 )));
396 }
397 }
398 }
399 _ => {
400 return Err(mlua::Error::RuntimeError(
401 "MAKE_ARCHIVE: source must be string or table".to_string(),
402 ));
403 }
404 }
405
406 let file = fs::File::create(&output_path)
407 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
408
409 match algo.as_str() {
410 "gz" => {
411 let mut encoder =
412 flate2::write::GzEncoder::new(file, flate2::Compression::default());
413 for (path, _) in source_paths {
414 let mut f = fs::File::open(path)
415 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
416 std::io::copy(&mut f, &mut encoder)
417 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
418 }
419 encoder
420 .finish()
421 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
422 }
423 "zip" => {
424 let mut zip = zip::ZipWriter::new(file);
425 let options = zip::write::SimpleFileOptions::default()
426 .compression_method(zip::CompressionMethod::Deflated);
427
428 for (path, rel_name) in source_paths {
429 if path.is_dir() {
430 for entry in walkdir::WalkDir::new(&path)
431 .into_iter()
432 .filter_map(|e| e.ok())
433 {
434 let rel =
435 entry.path().strip_prefix(path.parent().unwrap()).unwrap();
436 if entry.file_type().is_dir() {
437 zip.add_directory(rel.to_string_lossy(), options)
438 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
439 } else {
440 zip.start_file(rel.to_string_lossy(), options)
441 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
442 let mut f = fs::File::open(entry.path())
443 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
444 std::io::copy(&mut f, &mut zip)
445 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
446 }
447 }
448 } else {
449 zip.start_file(rel_name, options)
450 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
451 let mut f = fs::File::open(path)
452 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
453 std::io::copy(&mut f, &mut zip)
454 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
455 }
456 }
457 zip.finish()
458 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
459 }
460 "tar" | "tar.gz" | "tar.xz" | "tar.zst" | "zst" => {
461 let writer: Box<dyn std::io::Write> = match algo.as_str() {
462 "tar" => Box::new(file),
463 "tar.gz" => Box::new(flate2::write::GzEncoder::new(
464 file,
465 flate2::Compression::default(),
466 )),
467 "tar.xz" => Box::new(xz2::write::XzEncoder::new(file, 6)),
468 "tar.zst" | "zst" => Box::new(
469 zstd::stream::write::Encoder::new(file, 0)
470 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?
471 .auto_finish(),
472 ),
473 _ => unreachable!(),
474 };
475
476 let mut tar = tar::Builder::new(writer);
477 for (path, rel_name) in source_paths {
478 if path.is_dir() {
479 tar.append_dir_all(rel_name, path)
480 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
481 } else {
482 tar.append_path_with_name(path, rel_name)
483 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
484 }
485 }
486 tar.finish()
487 .map_err(|e| mlua::Error::RuntimeError(e.to_string()))?;
488 }
489 _ => {
490 return Err(mlua::Error::RuntimeError(format!(
491 "MAKE_ARCHIVE: unsupported algorithm: {}",
492 algo
493 )));
494 }
495 }
496
497 Ok(())
498 },
499 )?;
500
501 let utils_table: Table = lua.globals().get("UTILS")?;
502 utils_table.set("ARCHIVE", archive_table)?;
503 utils_table.set("MAKE_ARCHIVE", make_archive_fn)?;
504
505 Ok(())
506}