ecr/src/qemu_vm.rs

use crate::veprintln;
use anyhow::{anyhow, Context, Result};
use cpio::{newc, NewcBuilder};
use indicatif::{ProgressBar, ProgressStyle};
use std::io::Write;
use std::os::unix::fs::MetadataExt;
use std::path::{Path, PathBuf};
use std::process::Command;

/// QEMU system emulation configuration
pub struct QemuConfig {
    /// Path to the kernel image (vmlinuz)
    pub kernel_path: PathBuf,
    /// Path to the rootfs directory
    pub rootfs_path: PathBuf,
    /// Memory size for VM (e.g., "2G", "512M")
    pub memory: String,
    /// Target architecture
    pub arch: String,
    /// Optional command to run instead of default init
    pub command: Option<Vec<String>>,
}

/// Launch QEMU with the given configuration
pub fn launch_qemu(config: QemuConfig) -> Result<()> {
    // Check that kernel exists
    if !config.kernel_path.exists() {
        return Err(anyhow!(
            "Kernel not found: {}",
            config.kernel_path.display()
        ));
    }

    // Check that rootfs exists
    if !config.rootfs_path.exists() {
        return Err(anyhow!(
            "Rootfs not found: {}",
            config.rootfs_path.display()
        ));
    }

    // Create a gzipped cpio initramfs from the rootfs
    let initramfs = create_initramfs(&config.rootfs_path)?;

    // Get QEMU binary for architecture
    let qemu_bin = qemu_binary_for_arch(&config.arch);

    // Check QEMU exists
    which::which(&qemu_bin).context(format!(
        "QEMU system emulator '{}' not found. Install it with:\n\
         Ubuntu/Debian: sudo apt install qemu-system-{}\n\
         Arch: sudo pacman -S qemu-system-{}\n\
         Alpine: sudo apk add qemu-system-{}",
        qemu_bin, get_arch_package_suffix(&config.arch), get_arch_package_suffix(&config.arch), get_arch_package_suffix(&config.arch)
    ))?;

    // Detect the best available shell in the rootfs
    let shell = crate::utils::detect_shell(&config.rootfs_path);
    
    // Generate a unique hostname like "ecr-vm-a1b2c3"
    let hostname_suffix = format!("{:x}", (std::process::id() as u64)
        .wrapping_mul(std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_nanos() as u64) % 0x1000000);  // 6 hex chars
    let hostname = format!("ecr-vm-{}", hostname_suffix);
    
    // Build kernel command line
    // For initramfs boot, use rdinit= instead of init=
    // No root= needed as initramfs becomes the rootfs
    // 'quiet' suppresses kernel log messages for a cleaner console
    // Use setsid with redirected stdio to get proper job control
    // /dev/ttyS0 is created in the initramfs for serial console
    // Set hostname before spawning shell
    let kernel_append = if let Some(ref cmd) = config.command {
        let cmd_str = cmd.join(" ");
        format!(
            "console=ttyS0 quiet rdinit=/bin/sh -- -c \"echo {} >/etc/hostname; hostname {}; setsid sh -c 'exec {} </dev/ttyS0 >/dev/ttyS0 2>&1 -c {}'\"",
            hostname, hostname, shell, cmd_str
        )
    } else {
        // Default to interactive shell with proper job control
        format!(
            "console=ttyS0 quiet rdinit=/bin/sh -- -c \"echo {} >/etc/hostname; hostname {}; setsid sh -c 'exec {} </dev/ttyS0 >/dev/ttyS0 2>&1'\"",
            hostname, hostname, shell
        )
    };

    veprintln!("Launching QEMU: {}", qemu_bin);
    veprintln!("  Kernel: {}", config.kernel_path.display());
    veprintln!("  Initramfs: {}", initramfs.display());
    veprintln!("  Memory: {}", config.memory);
    veprintln!("  Kernel append: {}", kernel_append);

    // Build QEMU arguments
    // -display none suppresses VGA/BIOS output
    // -serial mon:stdio connects serial console to terminal with QEMU monitor muxed
    let args = vec![
        "-kernel".to_string(),
        config.kernel_path.to_string_lossy().to_string(),
        "-initrd".to_string(),
        initramfs.to_string_lossy().to_string(),
        "-append".to_string(),
        kernel_append,
        "-m".to_string(),
        config.memory.clone(),
        "-display".to_string(),
        "none".to_string(),
        "-serial".to_string(),
        "mon:stdio".to_string(),
        "-netdev".to_string(),
        "user,id=net0".to_string(),
        "-device".to_string(),
        "virtio-net-pci,netdev=net0".to_string(),
    ];

    // Execute QEMU
    let status = Command::new(&qemu_bin)
        .args(&args)
        .status()
        .context("Failed to execute QEMU")?;

    // Cleanup initramfs
    if let Err(e) = std::fs::remove_file(&initramfs) {
        veprintln!("Warning: failed to cleanup initramfs: {}", e);
    }

    if !status.success() {
        return Err(anyhow!(
            "QEMU exited with non-zero status: {}",
            status.code().unwrap_or(-1)
        ));
    }

    Ok(())
}

/// Get QEMU system binary name for architecture
fn qemu_binary_for_arch(arch: &str) -> String {
    match arch {
        "amd64" | "x86_64" => "qemu-system-x86_64".to_string(),
        "arm64" | "aarch64" => "qemu-system-aarch64".to_string(),
        "armhf" | "armv7" => "qemu-system-arm".to_string(),
        "riscv64" => "qemu-system-riscv64".to_string(),
        "ppc64el" | "ppc64le" => "qemu-system-ppc64".to_string(),
        "s390x" => "qemu-system-s390x".to_string(),
        other => format!("qemu-system-{}", other),
    }
}

/// Get architecture suffix for package names
fn get_arch_package_suffix(arch: &str) -> &str {
    match arch {
        "amd64" | "x86_64" => "x86",
        "arm64" | "aarch64" => "aarch64",
        "armhf" | "armv7" => "arm",
        "riscv64" => "riscv64",
        "ppc64el" | "ppc64le" => "ppc",
        "s390x" => "s390x",
        other => other,
    }
}

/// Create a gzipped cpio initramfs from a directory
fn create_initramfs(rootfs: &PathBuf) -> Result<PathBuf> {
    // Create a temporary file for the initramfs
    let initramfs_path = rootfs.parent().unwrap().join("initramfs.cpio.gz");

    // Create progress bar
    let pb = ProgressBar::new_spinner();
    pb.set_style(ProgressStyle::default_spinner()
        .template("{spinner:.green} {msg} ({pos} files)")
        .unwrap());
    pb.set_message("Scanning rootfs...");

    // Create the cpio archive with progress
    let cpio_data = create_cpio_archive(rootfs, &pb)?;
    let total_bytes = cpio_data.len() as u64;
    let file_count = pb.position();

    // Finish the scanning progress bar
    pb.finish_and_clear();

    // Create progress bar for compression
    let compress_pb = ProgressBar::new(total_bytes);
    compress_pb.set_style(ProgressStyle::default_bar()
        .template("{spinner:.green} Compressing initramfs... {bytes}/{total_bytes} ({eta})")
        .unwrap()
        .progress_chars("##-"));

    // Compress with gzip, writing in chunks to update progress
    let mut output_file = std::fs::File::create(&initramfs_path)
        .context("Failed to create initramfs file")?;
    
    let mut encoder = flate2::write::GzEncoder::new(&mut output_file, flate2::Compression::fast());
    
    // Write in 64KB chunks to provide progress updates
    const CHUNK_SIZE: usize = 64 * 1024;
    let mut offset = 0;
    while offset < cpio_data.len() {
        let end = std::cmp::min(offset + CHUNK_SIZE, cpio_data.len());
        encoder.write_all(&cpio_data[offset..end])
            .context("Failed to write compressed initramfs")?;
        offset = end;
        compress_pb.set_position(offset as u64);
    }
    encoder.finish()
        .context("Failed to finalize gzip compression")?;

    // Clear progress bar and print message only in verbose mode
    compress_pb.finish_and_clear();
    veprintln!("Initramfs created: {} bytes -> {} bytes compressed, {} files",
               total_bytes, output_file.metadata().ok().map(|m| m.len()).unwrap_or(0), file_count);

    Ok(initramfs_path)
}

/// Create a newc-format cpio archive from a directory using the cpio crate
fn create_cpio_archive(rootfs: &Path, pb: &ProgressBar) -> Result<Vec<u8>> {
    let mut archive = Vec::new();
    
    // Track seen inodes to handle hard links properly
    let mut seen_inodes = std::collections::HashMap::new();
    
    // Collect all entries with their data
    let entries = collect_entries(rootfs, rootfs, pb, &mut seen_inodes)?;
    
    // Collect entry names for checking existence later
    let entry_names: Vec<&str> = entries.iter().map(|(n, _, _, _, _)| n.as_str()).collect();
    
    pb.set_message("Writing initramfs...");
    
    // Write each entry using the cpio crate
    for (name, mode, mtime, nlink, data) in &entries {
        let file_size = data.len() as u32;
        
        let builder = NewcBuilder::new(name)
            .mode(*mode)
            .uid(0)
            .gid(0)
            .nlink(*nlink)
            .mtime(*mtime);
        
        // Write header and get a writer
        let mut writer = builder.write(&mut archive, file_size);
        
        // Write the file content
        writer.write_all(data)
            .context("Failed to write file content to cpio archive")?;
        
        // Finish this entry (returns the underlying writer)
        writer.finish()
            .context("Failed to finish cpio entry")?;
    }
    
    // Add essential device nodes for serial console
    // These are character devices (mode 0o020xxx)
    let device_nodes = [
        // /dev/ttyS0 - serial console (major 4, minor 64)
        ("dev/ttyS0", 0o020644, 4, 64),
        // /dev/null (major 1, minor 3)
        ("dev/null", 0o020644, 1, 3),
        // /dev/tty - controlling terminal (major 5, minor 0)
        ("dev/tty", 0o020666, 5, 0),
    ];
    
    for (name, mode, major, minor) in device_nodes {
        // Check if this device node already exists in the archive
        if entry_names.contains(&name) {
            continue;
        }
        
        let builder = NewcBuilder::new(name)
            .mode(mode)
            .uid(0)
            .gid(0)
            .nlink(1)
            .mtime(0)
            .rdev_major(major)
            .rdev_minor(minor);
        
        // Device nodes have zero size
        let writer = builder.write(&mut archive, 0);
        writer.finish()
            .context("Failed to finish device node entry")?;
    }
    
    // Write the trailer (takes ownership and returns the writer)
    archive = newc::trailer(archive)
        .context("Failed to write cpio trailer")?;

    Ok(archive)
}

/// Collect all filesystem entries recursively
/// Uses a HashMap to track hard links by (device, inode) - only stores data for first occurrence
fn collect_entries(
    base: &Path,
    current: &Path,
    pb: &ProgressBar,
    seen_inodes: &mut std::collections::HashMap<(u64, u64), String>,
) -> Result<Vec<(String, u32, u32, u32, Vec<u8>)>> {
    let mut entries = Vec::new();
    let mut total_data: u64 = 0;
    
    // Read directory entries
    let dir_entries: Vec<_> = match std::fs::read_dir(current) {
        Ok(entries) => entries.collect::<std::result::Result<_, _>>()?,
        Err(e) => {
            veprintln!("Warning: cannot read directory {}: {}", current.display(), e);
            return Ok(entries);
        }
    };
    
    for entry in dir_entries {
        let path = entry.path();
        
        // Get metadata
        let metadata = match std::fs::symlink_metadata(&path) {
            Ok(m) => m,
            Err(e) => {
                veprintln!("Warning: skipping {} due to metadata error: {}", path.display(), e);
                continue;
            }
        };
        
        // Increment progress counter
        pb.inc(1);
        
        let file_type = metadata.file_type();
        
        // Determine mode (file type + permissions from filesystem)
        let mode = if file_type.is_dir() {
            // Directory: preserve permissions, ensure at least rwx for owner
            0o040000 | (metadata.mode() & 0o7777)
        } else if file_type.is_symlink() {
            0o120777 // symlink with rwxrwxrwx (permissions don't matter for symlinks)
        } else if file_type.is_file() {
            // Regular file: preserve permissions from filesystem
            0o100000 | (metadata.mode() & 0o7777)
        } else {
            continue; // Skip other types (sockets, fifos, etc.)
        };
        
        // Build the entry name (relative path from base)
        let relative = path.strip_prefix(base).unwrap();
        let entry_name = relative.to_string_lossy().into_owned();
        
        // Handle hard links: only store data for first occurrence
        let (data, nlink) = if file_type.is_file() && metadata.nlink() > 1 {
            // This file has multiple hard links - check if we've seen it before
            let inode_key = (metadata.dev(), metadata.ino());
            
            if let Some(_first_path) = seen_inodes.get(&inode_key) {
                // We've seen this inode before - create a hard link entry with no data
                (Vec::new(), metadata.nlink() as u32)
            } else {
                // First occurrence - read the data and record this inode
                seen_inodes.insert(inode_key, entry_name.clone());
                let data = match std::fs::read(&path) {
                    Ok(data) => data,
                    Err(e) => {
                        veprintln!("Warning: cannot read file {}: {}", path.display(), e);
                        continue;
                    }
                };
                (data, metadata.nlink() as u32)
            }
        } else if file_type.is_file() {
            // Regular file with nlink=1
            let data = match std::fs::read(&path) {
                Ok(data) => data,
                Err(e) => {
                    veprintln!("Warning: cannot read file {}: {}", path.display(), e);
                    continue;
                }
            };
            (data, 1)
        } else if file_type.is_symlink() {
            match std::fs::read_link(&path) {
                Ok(target) => (target.to_string_lossy().into_owned().into_bytes(), 1),
                Err(e) => {
                    veprintln!("Warning: cannot read symlink {}: {}", path.display(), e);
                    continue;
                }
            }
        } else {
            // Directory
            (Vec::new(), 2)
        };
        
        total_data += data.len() as u64;
        entries.push((entry_name, mode, metadata.mtime() as u32, nlink, data));
        
        // Recurse into directories
        if file_type.is_dir() {
            let mut sub_entries = collect_entries(base, &path, pb, seen_inodes)?;
            for (_, _, _, _, d) in &sub_entries {
                total_data += d.len() as u64;
            }
            entries.append(&mut sub_entries);
        }
    }
    
    // Only print summary for the root directory
    if current == base {
        veprintln!("Collected {} entries, {} bytes total data", entries.len(), total_data);
    }
    Ok(entries)
}