use std::fs::{self, File};
use std::io::Write;
use std::path::Path;
use flate2::read::GzDecoder;
use flate2::write::GzEncoder;
use flate2::Compression;
use crate::error::{PwrError, Result};
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ArchiveStage {
Scanning,
Tarring,
Compressing,
Encrypting,
Hashing,
Decrypting,
Decompressing,
Extracting,
}
pub type ProgressFn = Box<dyn Fn(ArchiveStage, f64)>;
pub fn create_archive(
project_dir: &Path,
public_key: &str,
) -> Result<(Vec<u8>, String)> {
create_archive_with_progress(project_dir, public_key, None)
}
pub fn create_archive_with_progress(
project_dir: &Path,
public_key: &str,
progress: Option<&ProgressFn>,
) -> Result<(Vec<u8>, String)> {
if let Some(cb) = progress {
cb(ArchiveStage::Scanning, 0.0);
}
let tmp_dir = std::env::temp_dir().join(format!("pwr-archive-{}", uuid::Uuid::new_v4()));
fs::create_dir_all(&tmp_dir)?;
let tar_gz_path = tmp_dir.join("archive.tar.gz");
if let Some(cb) = progress {
cb(ArchiveStage::Tarring, 0.1);
}
let tar_gz_file = File::create(&tar_gz_path)?;
let encoder = GzEncoder::new(tar_gz_file, Compression::default());
let mut tar_builder = tar::Builder::new(encoder);
add_dir_to_tar(&mut tar_builder, project_dir, project_dir)?;
if let Some(cb) = progress {
cb(ArchiveStage::Compressing, 0.3);
}
let encoder = tar_builder
.into_inner()
.map_err(|e| PwrError::Crypto(format!("tar finalize error: {}", e)))?;
let _gz_file = encoder
.finish()
.map_err(|e| PwrError::Crypto(format!("gzip finalize error: {}", e)))?;
let tar_gz_data = fs::read(&tar_gz_path)?;
if let Some(cb) = progress {
cb(ArchiveStage::Encrypting, 0.5);
}
let encrypted = crate::crypto::age_encrypt(&tar_gz_data, public_key)?;
if let Some(cb) = progress {
cb(ArchiveStage::Hashing, 0.9);
}
let hash = crate::crypto::sha256_hex(&encrypted);
if let Some(cb) = progress {
cb(ArchiveStage::Hashing, 1.0);
}
let _ = fs::remove_dir_all(&tmp_dir);
Ok((encrypted, hash))
}
fn add_dir_to_tar<W: Write>(
builder: &mut tar::Builder<W>,
dir: &Path,
base_dir: &Path,
) -> Result<()> {
for entry in fs::read_dir(dir)? {
let entry = entry?;
let path = entry.path();
let relative = path
.strip_prefix(base_dir)
.map_err(|e| PwrError::Crypto(format!("path error: {}", e)))?;
if relative.as_os_str() == ".project.toml" {
continue;
}
if path.is_dir() {
builder
.append_dir(relative, &path)
.map_err(|e| PwrError::Crypto(format!("tar dir error: {}", e)))?;
add_dir_to_tar(builder, &path, base_dir)?;
} else if path.is_symlink() {
let target = fs::read_link(&path)?;
let mut header = tar::Header::new_gnu();
header.set_entry_type(tar::EntryType::Symlink);
header.set_size(0);
builder
.append_data(&mut header, relative, &mut &[][..])
.map_err(|e| PwrError::Crypto(format!("tar symlink error: {}", e)))?;
let _ = target;
log::debug!("Archived symlink: {}", relative.display());
} else {
builder
.append_path_with_name(&path, relative)
.map_err(|e| PwrError::Crypto(format!("tar file error: {}", e)))?;
}
}
Ok(())
}
pub fn extract_archive(
encrypted_blob: &[u8],
identity: &age::x25519::Identity,
target_dir: &Path,
expected_hash: &str,
) -> Result<()> {
extract_archive_with_progress(encrypted_blob, identity, target_dir, expected_hash, None)
}
pub fn extract_archive_with_progress(
encrypted_blob: &[u8],
identity: &age::x25519::Identity,
target_dir: &Path,
expected_hash: &str,
progress: Option<&ProgressFn>,
) -> Result<()> {
if let Some(cb) = progress {
cb(ArchiveStage::Hashing, 0.0);
}
let actual_hash = crate::crypto::sha256_hex(encrypted_blob);
if actual_hash != expected_hash {
return Err(PwrError::Crypto(format!(
"Hash mismatch: expected {}, got {}",
expected_hash, actual_hash
)));
}
if let Some(cb) = progress {
cb(ArchiveStage::Hashing, 0.2);
}
if let Some(cb) = progress {
cb(ArchiveStage::Decrypting, 0.3);
}
let decrypted = crate::crypto::age_decrypt(encrypted_blob, identity)?;
if let Some(cb) = progress {
cb(ArchiveStage::Decrypting, 0.5);
}
if let Some(cb) = progress {
cb(ArchiveStage::Decompressing, 0.6);
}
let gz_reader = GzDecoder::new(&decrypted[..]);
if let Some(cb) = progress {
cb(ArchiveStage::Extracting, 0.7);
}
fs::create_dir_all(target_dir)?;
let mut archive = tar::Archive::new(gz_reader);
archive
.unpack(target_dir)
.map_err(|e| PwrError::Crypto(format!("untar error: {}", e)))?;
if let Some(cb) = progress {
cb(ArchiveStage::Extracting, 1.0);
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
#[test]
fn test_create_and_extract_archive() -> Result<()> {
let tmp = TempDir::new()?;
let project_dir = tmp.path().join("testproj");
fs::create_dir_all(project_dir.join("src"))?;
fs::write(project_dir.join("README.md"), b"# Test Project\n")?;
fs::write(project_dir.join("src").join("main.rs"), b"fn main() {}\n")?;
fs::write(project_dir.join("Cargo.toml"), b"[package]\nname = \"test\"\n")?;
let identity = age::x25519::Identity::generate();
let public_key = identity.to_public().to_string();
let (encrypted, hash) = create_archive(&project_dir, &public_key)?;
assert!(!encrypted.is_empty());
assert!(!hash.is_empty());
let restore_dir = tmp.path().join("restored");
extract_archive(&encrypted, &identity, &restore_dir, &hash)?;
assert!(restore_dir.join("README.md").exists());
assert!(restore_dir.join("src").join("main.rs").exists());
assert!(restore_dir.join("Cargo.toml").exists());
assert!(!restore_dir.join(".project.toml").exists());
let readme = fs::read_to_string(restore_dir.join("README.md"))?;
assert_eq!(readme, "# Test Project\n");
Ok(())
}
#[test]
fn test_hash_verification_prevents_extraction() {
let tmp = TempDir::new().unwrap();
let project_dir = tmp.path().join("proj");
fs::create_dir_all(&project_dir).unwrap();
fs::write(project_dir.join("a.txt"), b"hello").unwrap();
let identity = age::x25519::Identity::generate();
let public_key = identity.to_public().to_string();
let (encrypted, _hash) = create_archive(&project_dir, &public_key).unwrap();
let result = extract_archive(
&encrypted,
&identity,
&tmp.path().join("bad"),
"0000000000000000000000000000000000000000000000000000000000000000",
);
assert!(result.is_err());
}
}