Skip to main content

zoi_lua/api/
archive.rs

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
12/// Exposes the `UTILS.EXTRACT` function to the Lua environment.
13///
14/// This utility provides a unified interface for downloading and extracting
15/// various archive formats. It handles:
16/// - Remote Fetching: If the source starts with http(s), it downloads the file to `BUILD_DIR`.
17/// - Format Detection: Dispatches to the appropriate decoder (Zip, Tar, Zstd, Xz, 7z, etc.).
18/// - Error Propagation: Any failure (network, filesystem, or corruption) is converted
19///   into an `mlua::Error::RuntimeError`, which halts the Lua execution and is
20///   caught by the Rust build engine to trigger a rollback.
21pub 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}