driftwood/src/core/repackager.rs

use std::fs;
use std::io::{Read, Write};
use std::path::{Path, PathBuf};
use super::database::Database;

/// Information about an AppImage's runtime binary.
#[derive(Debug, Clone)]
pub struct RuntimeInfo {
    pub runtime_size: u64,
    pub payload_offset: u64,
    pub runtime_type: RuntimeType,
    pub runtime_version: Option<String>,
}

/// The type of AppImage runtime.
#[derive(Debug, Clone, PartialEq)]
pub enum RuntimeType {
    OldFuse2,
    NewMulti,
    Static,
    Unknown,
}

impl RuntimeType {
    pub fn as_str(&self) -> &str {
        match self {
            Self::OldFuse2 => "old-fuse2",
            Self::NewMulti => "new-multi",
            Self::Static => "static",
            Self::Unknown => "unknown",
        }
    }

    pub fn label(&self) -> &str {
        match self {
            Self::OldFuse2 => "Legacy FUSE 2 only",
            Self::NewMulti => "Multi-runtime (FUSE 2/3 + static)",
            Self::Static => "Static (no FUSE needed)",
            Self::Unknown => "Unknown runtime",
        }
    }
}

/// Result of a runtime replacement operation.
#[derive(Debug)]
pub struct RepackageResult {
    pub original_path: PathBuf,
    pub backup_path: PathBuf,
    pub old_runtime_type: RuntimeType,
    pub new_runtime_type: String,
    pub old_size: u64,
    pub new_size: u64,
    pub success: bool,
}

/// Detect the runtime type and payload offset of an AppImage.
/// Type 2 AppImages store the SquashFS offset in the ELF section header.
pub fn detect_runtime(appimage_path: &Path) -> Result<RuntimeInfo, RepackageError> {
    let mut file = fs::File::open(appimage_path)
        .map_err(|e| RepackageError::Io(e.to_string()))?;

    // Read ELF header to find section headers
    let mut header = [0u8; 64];
    file.read_exact(&mut header)
        .map_err(|e| RepackageError::Io(e.to_string()))?;

    // Verify ELF magic
    if &header[0..4] != b"\x7fELF" {
        return Err(RepackageError::NotAppImage("Not an ELF file".to_string()));
    }

    // Find the SquashFS payload by searching for the magic bytes
    let payload_offset = find_squashfs_offset(appimage_path)?;

    let runtime_size = payload_offset;

    // Classify the runtime type based on size and content
    let runtime_type = classify_runtime(appimage_path, runtime_size)?;

    Ok(RuntimeInfo {
        runtime_size,
        payload_offset,
        runtime_type,
        runtime_version: None,
    })
}

/// Find the offset where the SquashFS payload starts.
/// SquashFS magic is 'hsqs' (0x73717368) at the start of the payload.
fn find_squashfs_offset(appimage_path: &Path) -> Result<u64, RepackageError> {
    let mut file = fs::File::open(appimage_path)
        .map_err(|e| RepackageError::Io(e.to_string()))?;

    let file_size = file.metadata()
        .map(|m| m.len())
        .map_err(|e| RepackageError::Io(e.to_string()))?;

    // SquashFS magic: 'hsqs' = [0x68, 0x73, 0x71, 0x73]
    let magic = b"hsqs";

    // Search in chunks starting from reasonable offsets (runtime is typically 100-300KB)
    let mut buf = [0u8; 65536];
    let search_start = 4096u64; // Skip the ELF header
    let search_end = std::cmp::min(file_size, 1_048_576); // Don't search beyond 1MB

    let mut offset = search_start;
    use std::io::Seek;
    file.seek(std::io::SeekFrom::Start(offset))
        .map_err(|e| RepackageError::Io(e.to_string()))?;

    while offset < search_end {
        let n = file.read(&mut buf)
            .map_err(|e| RepackageError::Io(e.to_string()))?;
        if n == 0 { break; }

        // Search for magic in this chunk
        for i in 0..n.saturating_sub(3) {
            if &buf[i..i + 4] == magic {
                return Ok(offset + i as u64);
            }
        }

        offset += n as u64 - 3; // Overlap by 3 to catch magic spanning chunks
        file.seek(std::io::SeekFrom::Start(offset))
            .map_err(|e| RepackageError::Io(e.to_string()))?;
    }

    Err(RepackageError::NotAppImage("SquashFS payload not found".to_string()))
}

/// Classify the runtime type based on its binary content.
fn classify_runtime(appimage_path: &Path, runtime_size: u64) -> Result<RuntimeType, RepackageError> {
    let mut file = fs::File::open(appimage_path)
        .map_err(|e| RepackageError::Io(e.to_string()))?;

    let read_size = std::cmp::min(runtime_size, 65536) as usize;
    let mut buf = vec![0u8; read_size];
    file.read_exact(&mut buf)
        .map_err(|e| RepackageError::Io(e.to_string()))?;

    let content = String::from_utf8_lossy(&buf);

    // Check for known strings in the runtime binary
    if content.contains("libfuse3") || content.contains("fuse3") {
        Ok(RuntimeType::NewMulti)
    } else if content.contains("static-runtime") || content.contains("no-fuse") {
        Ok(RuntimeType::Static)
    } else if content.contains("libfuse") || content.contains("fuse2") {
        Ok(RuntimeType::OldFuse2)
    } else if runtime_size < 4096 {
        // Suspiciously small runtime - probably not a valid AppImage runtime
        Ok(RuntimeType::Unknown)
    } else {
        // Default: older runtimes are typically fuse2-only
        Ok(RuntimeType::OldFuse2)
    }
}

/// Replace the runtime of an AppImage with a new one.
/// Creates a backup of the original file before modifying.
pub fn replace_runtime(
    appimage_path: &Path,
    new_runtime_path: &Path,
    keep_backup: bool,
) -> Result<RepackageResult, RepackageError> {
    if !appimage_path.exists() {
        return Err(RepackageError::NotAppImage("File not found".to_string()));
    }
    if !new_runtime_path.exists() {
        return Err(RepackageError::Io("New runtime file not found".to_string()));
    }

    let info = detect_runtime(appimage_path)?;
    let old_size = fs::metadata(appimage_path)
        .map(|m| m.len())
        .map_err(|e| RepackageError::Io(e.to_string()))?;

    // Create backup
    let backup_path = appimage_path.with_extension("bak");
    fs::copy(appimage_path, &backup_path)
        .map_err(|e| RepackageError::Io(format!("Backup failed: {}", e)))?;

    // Read new runtime
    let new_runtime = fs::read(new_runtime_path)
        .map_err(|e| RepackageError::Io(format!("Failed to read new runtime: {}", e)))?;

    // Read the SquashFS payload from the original file
    let mut original = fs::File::open(appimage_path)
        .map_err(|e| RepackageError::Io(e.to_string()))?;
    use std::io::Seek;
    original.seek(std::io::SeekFrom::Start(info.payload_offset))
        .map_err(|e| RepackageError::Io(e.to_string()))?;

    let mut payload = Vec::new();
    original.read_to_end(&mut payload)
        .map_err(|e| RepackageError::Io(e.to_string()))?;
    drop(original);

    // Write new AppImage: new_runtime + payload
    let mut output = fs::File::create(appimage_path)
        .map_err(|e| RepackageError::Io(e.to_string()))?;
    output.write_all(&new_runtime)
        .map_err(|e| RepackageError::Io(e.to_string()))?;
    output.write_all(&payload)
        .map_err(|e| RepackageError::Io(e.to_string()))?;

    // Set executable permission
    #[cfg(unix)]
    {
        use std::os::unix::fs::PermissionsExt;
        let perms = fs::Permissions::from_mode(0o755);
        fs::set_permissions(appimage_path, perms).ok();
    }

    let new_size = fs::metadata(appimage_path)
        .map(|m| m.len())
        .unwrap_or(0);

    // Verify the new file is a valid AppImage
    let success = verify_appimage(appimage_path);

    if !success {
        // Rollback from backup
        log::error!("Verification failed, rolling back from backup");
        fs::copy(&backup_path, appimage_path).ok();
        if !keep_backup {
            fs::remove_file(&backup_path).ok();
        }
        return Err(RepackageError::VerificationFailed);
    }

    if !keep_backup {
        fs::remove_file(&backup_path).ok();
    }

    Ok(RepackageResult {
        original_path: appimage_path.to_path_buf(),
        backup_path,
        old_runtime_type: info.runtime_type,
        new_runtime_type: "new".to_string(),
        old_size,
        new_size,
        success: true,
    })
}

/// Batch-replace runtimes for all AppImages in the database that use the old runtime.
pub fn batch_replace_runtimes(
    db: &Database,
    new_runtime_path: &Path,
    dry_run: bool,
) -> Vec<RepackageResult> {
    let records = db.get_all_appimages().unwrap_or_default();
    let mut results = Vec::new();

    for record in &records {
        let path = Path::new(&record.path);
        if !path.exists() {
            continue;
        }

        let info = match detect_runtime(path) {
            Ok(i) => i,
            Err(e) => {
                log::warn!("Skipping {}: {}", record.filename, e);
                continue;
            }
        };

        // Only repackage old fuse2 runtimes
        if info.runtime_type != RuntimeType::OldFuse2 {
            continue;
        }

        if dry_run {
            results.push(RepackageResult {
                original_path: path.to_path_buf(),
                backup_path: path.with_extension("bak"),
                old_runtime_type: info.runtime_type,
                new_runtime_type: "new".to_string(),
                old_size: fs::metadata(path).map(|m| m.len()).unwrap_or(0),
                new_size: 0,
                success: true,
            });
            continue;
        }

        match replace_runtime(path, new_runtime_path, true) {
            Ok(result) => {
                // Record in database
                db.record_runtime_update(
                    record.id,
                    Some(info.runtime_type.as_str()),
                    Some("new"),
                    result.backup_path.to_str(),
                    true,
                ).ok();
                results.push(result);
            }
            Err(e) => {
                log::error!("Failed to repackage {}: {}", record.filename, e);
                db.record_runtime_update(
                    record.id,
                    Some(info.runtime_type.as_str()),
                    Some("new"),
                    None,
                    false,
                ).ok();
            }
        }
    }

    results
}

/// Download the latest AppImage runtime binary.
pub fn download_latest_runtime() -> Result<PathBuf, RepackageError> {
    let url = "https://github.com/AppImage/type2-runtime/releases/latest/download/runtime-x86_64";

    let dest = dirs::cache_dir()
        .unwrap_or_else(|| PathBuf::from("/tmp"))
        .join("driftwood")
        .join("runtime-x86_64");

    fs::create_dir_all(dest.parent().unwrap()).ok();

    let response = ureq::get(url)
        .call()
        .map_err(|e| RepackageError::Network(e.to_string()))?;

    let mut file = fs::File::create(&dest)
        .map_err(|e| RepackageError::Io(e.to_string()))?;

    let mut reader = response.into_body().into_reader();
    let mut buf = [0u8; 65536];
    loop {
        let n = reader.read(&mut buf)
            .map_err(|e| RepackageError::Network(e.to_string()))?;
        if n == 0 { break; }
        file.write_all(&buf[..n])
            .map_err(|e| RepackageError::Io(e.to_string()))?;
    }

    #[cfg(unix)]
    {
        use std::os::unix::fs::PermissionsExt;
        fs::set_permissions(&dest, fs::Permissions::from_mode(0o755)).ok();
    }

    Ok(dest)
}

/// Basic verification that a file is still a valid AppImage.
fn verify_appimage(path: &Path) -> bool {
    // Check ELF magic
    let mut file = match fs::File::open(path) {
        Ok(f) => f,
        Err(_) => return false,
    };

    let mut magic = [0u8; 4];
    if file.read_exact(&mut magic).is_err() {
        return false;
    }
    if &magic != b"\x7fELF" {
        return false;
    }

    // Check that SquashFS payload exists
    find_squashfs_offset(path).is_ok()
}

// --- Error types ---

#[derive(Debug)]
pub enum RepackageError {
    NotAppImage(String),
    Io(String),
    Network(String),
    VerificationFailed,
}

impl std::fmt::Display for RepackageError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::NotAppImage(e) => write!(f, "Not a valid AppImage: {}", e),
            Self::Io(e) => write!(f, "I/O error: {}", e),
            Self::Network(e) => write!(f, "Network error: {}", e),
            Self::VerificationFailed => write!(f, "Verification failed after repackaging"),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_runtime_type_as_str() {
        assert_eq!(RuntimeType::OldFuse2.as_str(), "old-fuse2");
        assert_eq!(RuntimeType::NewMulti.as_str(), "new-multi");
        assert_eq!(RuntimeType::Static.as_str(), "static");
        assert_eq!(RuntimeType::Unknown.as_str(), "unknown");
    }

    #[test]
    fn test_runtime_type_label() {
        assert!(RuntimeType::OldFuse2.label().contains("Legacy"));
        assert!(RuntimeType::NewMulti.label().contains("Multi"));
        assert!(RuntimeType::Static.label().contains("no FUSE"));
    }

    #[test]
    fn test_repackage_error_display() {
        let err = RepackageError::NotAppImage("bad magic".to_string());
        assert!(format!("{}", err).contains("bad magic"));
        let err = RepackageError::VerificationFailed;
        assert!(format!("{}", err).contains("Verification failed"));
    }

    #[test]
    fn test_detect_runtime_nonexistent() {
        let result = detect_runtime(Path::new("/nonexistent.AppImage"));
        assert!(result.is_err());
    }

    #[test]
    fn test_detect_runtime_not_elf() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("not-an-elf");
        fs::write(&path, "This is not an ELF file").unwrap();
        let result = detect_runtime(&path);
        assert!(result.is_err());
    }

    #[test]
    fn test_verify_appimage_nonexistent() {
        assert!(!verify_appimage(Path::new("/nonexistent")));
    }

    #[test]
    fn test_verify_appimage_not_elf() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("not-elf");
        fs::write(&path, "hello").unwrap();
        assert!(!verify_appimage(&path));
    }
}