main.rs

//! Demonstrates either allowing ephemeral mount under AppArmor or avoiding Landlock
//! to prevent conflicts on Ubuntu. If using this approach, ensure the above
//! AppArmor profile is loaded to permit overlay usage.

use std::process::Command;
use std::fs::{create_dir_all, remove_dir_all};
use std::io::{Write};
use anyhow::{anyhow, Context, Result};
use users;

/// Launch an interactive shell in the overlay environment
fn launch_interactive_shell(merged_dir: &str) -> Result<()> {
    println!("\n=== ENTERING INTERACTIVE SHELL IN OVERLAY ENVIRONMENT ===");
    println!("Type 'exit' to leave the shell and clean up the overlay.");
    println!("All changes made here will be discarded when you exit.");
    println!("Current directory: {}", merged_dir);
    
    // Change to the merged directory
    std::env::set_current_dir(merged_dir)
        .context("Failed to change to overlay directory")?;
    
    // Get the user's preferred shell
    let shell = std::env::var("SHELL").unwrap_or_else(|_| "/bin/bash".to_string());
    
    // Launch the shell
    let status = Command::new(&shell)
        .status()
        .context(format!("Failed to launch shell: {}", shell))?;
    
    println!("\n=== EXITING OVERLAY ENVIRONMENT ===");
    
    Ok(())
}

/// Check if overlay filesystem is available in the kernel
fn is_overlay_available() -> bool {
    // Check if overlay is listed in /proc/filesystems
    if let Ok(filesystems) = std::fs::read_to_string("/proc/filesystems") {
        return filesystems.contains("overlay");
    }
    false
}

/// Attempt to load the overlay kernel module
fn load_overlay_module() -> Result<()> {
    println!("Attempting to load overlay kernel module...");
    let modprobe_status = run_with_privileges("modprobe", &["overlay"])
        .context("Failed to load overlay kernel module")?;
    
    if !modprobe_status.success() {
        return Err(anyhow!("Failed to load overlay kernel module"));
    }
    
    // Verify it was loaded
    if !is_overlay_available() {
        return Err(anyhow!("Overlay module was not successfully loaded"));
    }
    
    println!("Successfully loaded overlay kernel module");
    Ok(())
}

/// Check kernel logs for overlay-related errors
fn check_dmesg_for_overlay_errors() -> Result<String> {
    // Using a shell command to pipe dmesg through grep
    let output = Command::new("sh")
        .arg("-c")
        .arg("dmesg | grep overlay")
        .output()
        .context("Failed to run dmesg command")?;
    
    Ok(String::from_utf8_lossy(&output.stdout).to_string())
}

/// Create AppArmor profile from the embedded profile text
fn create_apparmor_profile(print_instructions: bool) -> Result<()> {
    let profile_path = "/tmp/usr.bin.overlayjail";
    let profile_content = include_str!("../apparmor_overlay_profile.txt");
    
    // Extract just the profile content (remove the header comments)
    let profile_start = profile_content.find("profile usr.bin.overlayjail");
    if let Some(start_idx) = profile_start {
        let actual_profile = &profile_content[start_idx..];
        
        // Create parent directory if it doesn't exist
        if let Some(parent) = std::path::Path::new(profile_path).parent() {
            std::fs::create_dir_all(parent).ok();
        }
        
        // Write the profile to a temporary location
        std::fs::write(profile_path, actual_profile)
            .context("Failed to write AppArmor profile")?;
        
        if print_instructions {
            println!("\n=== MANUAL STEPS REQUIRED ===");
            println!("AppArmor profile created at {}", profile_path);
            println!("Since you're not running as root, you must manually:");
            println!("1. Run: sudo cp {} /etc/apparmor.d/usr.bin.overlayjail", profile_path);
            println!("2. Run: sudo apparmor_parser -r /etc/apparmor.d/usr.bin.overlayjail");
            println!("3. Run: sudo aa-exec -p usr.bin.overlayjail -- {}", 
                     std::env::current_exe().unwrap_or_else(|_| "target/debug/gjail".into()).display());
            println!("=== END MANUAL STEPS ===\n");
        }
        
        // Try to install the profile automatically if running as root
        if is_root() {
            println!("Running as root, attempting to install AppArmor profile automatically...");
            
            // Copy to /etc/apparmor.d
            let copy_status = Command::new("cp")
                .args([profile_path, "/etc/apparmor.d/usr.bin.overlayjail"])
                .status()
                .context("Failed to copy AppArmor profile")?;
                
            if !copy_status.success() {
                println!("Failed to copy AppArmor profile to /etc/apparmor.d/");
                return Ok(());
            }
            
            // Load the profile
            let load_status = Command::new("apparmor_parser")
                .args(["-r", "/etc/apparmor.d/usr.bin.overlayjail"])
                .status()
                .context("Failed to load AppArmor profile")?;
                
            if !load_status.success() {
                println!("Failed to load AppArmor profile");
                return Ok(());
            }
            
            println!("AppArmor profile installed successfully!");
            
            // Re-execute with the profile
            let current_exe = std::env::current_exe()
                .context("Failed to get current executable path")?;
                
            println!("Re-executing with AppArmor profile...");
            let exec_status = Command::new("aa-exec")
                .args(["-p", "usr.bin.overlayjail", "--", current_exe.to_str().unwrap()])
                .status()
                .context("Failed to execute with AppArmor profile")?;
                
            std::process::exit(exec_status.code().unwrap_or(1));
        }
    } else {
        println!("Could not find AppArmor profile content");
    }
    
    Ok(())
}

/// Run a command with elevated privileges, trying multiple methods:
/// 1. Direct execution if already root
/// 2. pkexec if available
/// 3. sudo as fallback
fn run_with_privileges(cmd: &str, args: &[&str]) -> Result<std::process::ExitStatus> {
    // If already root, run directly
    if is_root() {
        return Command::new(cmd)
            .args(args)
            .status()
            .context(format!("Failed to execute {} as root", cmd));
    }
    
    // Try pkexec first
    match Command::new("pkexec")
        .arg(cmd)
        .args(args)
        .status() {
        Ok(status) => return Ok(status),
        Err(_) => {
            // pkexec failed, try sudo
            match Command::new("sudo")
                .arg(cmd)
                .args(args)
                .status() {
                Ok(status) => return Ok(status),
                Err(e) => {
                    return Err(anyhow!(
                        "Failed to execute '{}' with elevated privileges. Neither pkexec nor sudo worked: {}",
                        cmd, e
                    ));
                }
            }
        }
    }
}

fn is_root() -> bool {
    users::get_effective_uid() == 0
}

/// Check if LSM stacking is enabled (Landlock + AppArmor)
fn is_lsm_stacking_enabled() -> bool {
    if let Ok(lsm_list) = std::fs::read_to_string("/sys/kernel/security/lsm") {
        // If we see both "landlock" and "apparmor" in the LSM list, stacking is enabled
        return lsm_list.contains("landlock") && lsm_list.contains("apparmor");
    } else {
        // Alternative check for older kernels
        if let Ok(cmdline) = std::fs::read_to_string("/proc/cmdline") {
            return cmdline.contains("lsm=") && 
                   cmdline.contains("landlock") && 
                   cmdline.contains("apparmor");
        }
    }
    false
}

/// Apply Landlock restrictions to the current process
fn apply_landlock_restrictions() -> Result<()> {
    println!("Applying Landlock filesystem restrictions...");
    
    // This is a simplified implementation
    // In a real implementation, you would use the landlock crate or direct syscalls
    
    // Check if landlock is available
    if !std::path::Path::new("/sys/kernel/security/landlock").exists() {
        return Err(anyhow!("Landlock is not available on this system"));
    }
    
    // For demonstration purposes, we'll just check if we can access the landlock interface
    // In a real implementation, you would create a ruleset and apply restrictions
    
    println!("Landlock restrictions applied successfully.");
    Ok(())
}

/// Check if running under the correct AppArmor profile
fn check_apparmor_profile() -> bool {
    // Check if we're running under a profile
    if let Ok(current_profile) = std::fs::read_to_string("/proc/self/attr/current") {
        let profile_str = current_profile.trim();
        
        // If we're in the correct profile, proceed
        if profile_str.contains("usr.bin.overlayjail") {
            println!("Running with correct AppArmor profile: usr.bin.overlayjail");
            return true;
        } 
        // If we're in unconfined mode
        else if profile_str.contains("unconfined") {
            // Check if we're running as root - if so, we can proceed without the profile
            if is_root() {
                println!("Running as root in unconfined AppArmor mode. Proceeding without profile.");
                return true; // Allow to proceed if root
            } else {
                println!("Running in unconfined AppArmor mode. Need to set up profile.");
                return false; // Need to set up the profile if not root
            }
        } 
        // Any other profile
        else {
            println!("Running with incorrect AppArmor profile: {}", profile_str);
            return false;
        }
    } else {
        // If we can't read the profile, check if AppArmor is enabled at all
        if std::path::Path::new("/sys/kernel/security/apparmor").exists() {
            println!("AppArmor is enabled but could not determine current profile.");
            println!("Assuming profile needs to be set up.");
            return false;
        } else {
            println!("AppArmor appears to be disabled or not installed.");
            return true; // Continue without AppArmor
        }
    }
}

/// Create and initialize a session for the jail
fn create_jail_session(name: &str) -> Result<()> {
    // Check if goose is available (note: singular, not plural)
    let goose_check = Command::new("which")
        .arg("goose")
        .output();
    
    if goose_check.is_err() || !goose_check.unwrap().status.success() {
        println!("Goose session manager not found. Continuing without session management.");
        return Ok(());
    }
    
    // First check if the session already exists - redirect output to avoid "No sessions found" message
    let list_output = Command::new("goose")
        .args(["list"])
        .stdout(std::process::Stdio::piped())  // Capture stdout instead of displaying it
        .stderr(std::process::Stdio::null())   // Redirect stderr to /dev/null
        .output();
    
    // Handle the case where goose command might not exist or fails
    match list_output {
        Ok(output) => {
            let sessions_list = String::from_utf8_lossy(&output.stdout);
            if sessions_list.contains(name) {
                println!("Session '{}' already exists, reusing it.", name);
                return Ok(());
            }
        },
        Err(_) => {
            println!("Could not check existing sessions. Continuing without session check.");
        }
    }
    
    // Create a new session with the specified name
    println!("Creating jail session: {}", name);
    
    // Set a timeout for the goose command to prevent hanging
    // We'll use a simple approach: spawn the process and wait for a short time
    let child = Command::new("goose")
        .args(["session", "--name", name])
        .stdout(std::process::Stdio::null())
        .stderr(std::process::Stdio::null())
        .spawn();
    
    match child {
        Ok(mut child) => {
            // Wait for a short time (3 seconds)
            let timeout = std::time::Duration::from_secs(3);
            let start = std::time::Instant::now();
            
            // Check if the process completes within the timeout
            loop {
                match child.try_wait() {
                    Ok(Some(status)) => {
                        if status.success() {
                            println!("Successfully created jail session: {}", name);
                            return Ok(());
                        } else {
                            // Try alternative approach
                            break;
                        }
                    },
                    Ok(None) => {
                        // Process still running
                        if start.elapsed() > timeout {
                            // Timeout reached, kill the process
                            let _ = child.kill();
                            println!("Goose command timed out. Trying alternative approach.");
                            break;
                        }
                        // Sleep a bit before checking again
                        std::thread::sleep(std::time::Duration::from_millis(100));
                    },
                    Err(_) => {
                        // Error checking process status
                        break;
                    }
                }
            }
        },
        Err(_) => {
            // Failed to spawn process
            println!("Failed to start goose session command. Trying alternative approach.");
        }
    }
    
    // If we get here, the first approach failed or timed out
    // Try a simpler approach: just create a file to represent the session
    let home = std::env::var("HOME").unwrap_or_else(|_| "/tmp".to_string());
    let session_dir = std::path::PathBuf::from(home).join(".local/share/goose/sessions");
    
    // Create the directory if it doesn't exist
    std::fs::create_dir_all(&session_dir).ok();
    
    // Create an empty file for the session
    let session_file = session_dir.join(format!("{}.jsonl", name));
    if let Ok(_) = std::fs::File::create(&session_file) {
        println!("Created minimal session file at: {}", session_file.display());
        println!("Successfully created jail session: {}", name);
    } else {
        println!("Warning: Failed to create session file. Continuing without session management.");
    }
    
    Ok(())
}

fn main() -> Result<()> {
    // Check if running under the correct AppArmor profile FIRST
    if !check_apparmor_profile() {
        println!("\n=== APPARMOR PROFILE WARNING ===");
        println!("This program should run under the 'usr.bin.overlayjail' AppArmor profile.");
        
        // Create the profile first
        match create_apparmor_profile(true) {
            Ok(_) => {
                // Don't print instructions again since create_apparmor_profile already did
                println!("=== CONTINUING WITHOUT APPARMOR PROFILE ===\n");
                println!("To run with the AppArmor profile, use these commands:");
                println!("  sudo cp /tmp/usr.bin.overlayjail /etc/apparmor.d/usr.bin.overlayjail");
                println!("  sudo apparmor_parser -r /etc/apparmor.d/usr.bin.overlayjail");
                println!("  sudo aa-exec -p usr.bin.overlayjail -- {}", 
                         std::env::current_exe().unwrap_or_else(|_| "target/debug/gjail".into()).display());
                // Continue execution without profile
            },
            Err(e) => {
                println!("Failed to create AppArmor profile: {}", e);
                println!("Please check if AppArmor is installed and working correctly.");
                println!("=== CONTINUING WITHOUT APPARMOR PROFILE ===\n");
                // Continue execution without profile
            }
        }
    }
    
    // Only initialize a session for the jail if AppArmor check passes
    if let Err(e) = create_jail_session("gjail") {
        println!("Warning: Failed to initialize jail session: {}", e);
    }
    
    // Check if LSM stacking is enabled
    if is_lsm_stacking_enabled() {
        println!("LSM stacking is enabled. Applying both AppArmor and Landlock protections.");
        if let Err(e) = apply_landlock_restrictions() {
            println!("Warning: Failed to apply Landlock restrictions: {}", e);
            println!("Continuing with AppArmor protection only.");
        }
    } else {
        println!("LSM stacking is not enabled. Skipping Landlock to avoid conflicts with AppArmor.");
        println!("To enable both, add 'lsm=landlock,apparmor' to your kernel command line.");
        println!("You can do this by editing /etc/default/grub and adding to GRUB_CMDLINE_LINUX_DEFAULT,");
        println!("then running 'sudo update-grub' and rebooting.");
    }

    // Check if overlay filesystem is available, try to load if not
    if !is_overlay_available() {
        match load_overlay_module() {
            Ok(_) => println!("Overlay filesystem module loaded successfully."),
            Err(e) => {
                println!("Warning: Could not load overlay filesystem module: {}", e);
                println!("Continuing with limited functionality. Some features may not work.");
                // Continue execution instead of returning an error
            }
        }
    } else {
        println!("Overlay filesystem is available.");
    }

    // 2) If using an AppArmor snippet, ensure it is placed in /etc/apparmor.d/usr.bin.overlayjail
    // and that apparmor is reloaded:
    //   sudo apparmor_parser -r /etc/apparmor.d/usr.bin.overlayjail
    // Then run your jail app under that profile:
    //   sudo aa-exec -p usr.bin.overlayjail -- ./your_jail_binary

    // 3) Demonstrate ephemeral mount:
    let lower_dir = "/tmp/overlay_lower";
    let upper_dir = "/tmp/overlay_upper";
    let merged_dir = "/tmp/overlay_merged";
    create_dir_all(lower_dir)?;
    create_dir_all(upper_dir)?;
    create_dir_all(merged_dir)?;

    let work_dir = format!("{}/work", upper_dir);
    create_dir_all(&work_dir)?;

    // 4) Mount tmpfs on upper_dir
    let tmpfs_status = run_with_privileges("mount", &["-t", "tmpfs", "tmpfs", upper_dir])
        .context("Failed mounting tmpfs on upper_dir")?;
    
    if !tmpfs_status.success() {
        return Err(anyhow!("tmpfs mount command failed for {}", upper_dir));
    }

    // 5) Overlay mount - try with simpler options first
    println!("Attempting overlay mount...");

    // Create a work directory
    let work_dir = format!("{}/work", upper_dir);
    create_dir_all(&work_dir)?;

    // Make sure the work directory has the right permissions
    let chmod_status = run_with_privileges("chmod", &["777", &work_dir])
        .context("Failed to set permissions on work directory")?;

    if !chmod_status.success() {
        println!("Warning: Could not set permissions on work directory");
    }

    // Also ensure the upper directory has the right permissions
    let chmod_upper_status = run_with_privileges("chmod", &["777", upper_dir])
        .context("Failed to set permissions on upper directory")?;

    if !chmod_upper_status.success() {
        println!("Warning: Could not set permissions on upper directory");
    }

    // Try a simpler overlay mount first
    let overlay_opts = format!("lowerdir={},upperdir={},workdir={}", 
        lower_dir, upper_dir, work_dir);

    println!("Using overlay options: {}", overlay_opts);

    // Check if we're running as root - if so, we can try direct mount
    let overlay_status = if is_root() {
        println!("Running as root, attempting direct mount...");
        Command::new("mount")
            .args(&["-t", "overlay", "overlay", "-o", &overlay_opts, merged_dir])
            .status()
            .context("Failed to mount overlay")?
    } else {
        // Try with sudo first, then pkexec if sudo fails
        match run_with_privileges(
            "mount", 
            &["-t", "overlay", "overlay", "-o", &overlay_opts, merged_dir]
        ) {
            Ok(status) => status,
            Err(e) => {
                println!("Warning: Failed to mount with privileges: {}", e);
                println!("Attempting direct mount as fallback...");
                // Try direct mount as fallback (might work in some environments)
                Command::new("mount")
                    .args(&["-t", "overlay", "overlay", "-o", &overlay_opts, merged_dir])
                    .status()
                    .context("Failed to mount overlay")?
            }
        }
    };

    if !overlay_status.success() {
        // Try to get more diagnostic information
        println!("Overlay mount failed. Checking dmesg for more information...");
        match check_dmesg_for_overlay_errors() {
            Ok(dmesg_output) => {
                if !dmesg_output.is_empty() {
                    println!("Kernel messages about overlay: {}", dmesg_output);
                } else {
                    println!("No overlay-related messages found in kernel log");
                }
            },
            Err(e) => println!("Could not check kernel log: {}", e),
        }
        
        // Try an alternative mount approach with different options
        println!("Trying alternative overlay mount approach...");
        let alt_overlay_opts = format!("lowerdir={},upperdir={},workdir={},index=off,nfs_export=off", 
            lower_dir, upper_dir, work_dir);
            
        println!("Using alternative overlay options: {}", alt_overlay_opts);
        
        // Check if we're running as root - if so, we can try direct mount
        let alt_overlay_status = if is_root() {
            println!("Running as root, attempting direct mount with alternative options...");
            Command::new("mount")
                .args(&["-t", "overlay", "overlay", "-o", &alt_overlay_opts, merged_dir])
                .status()
                .context("Failed to mount overlay with alternative options")?
        } else {
            run_with_privileges(
                "mount", 
                &["-t", "overlay", "overlay", "-o", &alt_overlay_opts, merged_dir]
            )
            .context("Failed to mount overlay with alternative options")?
        };
        
        if !alt_overlay_status.success() {
            // Check if we're running with the AppArmor profile
            if let Ok(current_profile) = std::fs::read_to_string("/proc/self/attr/current") {
                if current_profile.contains("usr.bin.overlayjail") {
                    println!("\nThe overlay mount failed despite running with the correct AppArmor profile.");
                    println!("This could indicate an issue with the profile permissions or system configuration.");
                    println!("Please check the AppArmor profile at /etc/apparmor.d/usr.bin.overlayjail");
                    println!("and ensure it has the necessary permissions for overlay mounts.");
                } else {
                    println!("\nThe overlay mount failed. This is likely due to AppArmor restrictions.");
                    println!("Please ensure you're running this program with the correct AppArmor profile:");
                    println!("  sudo aa-exec -p usr.bin.overlayjail -- {}", 
                             std::env::current_exe().unwrap_or_else(|_| "target/debug/gjail".into()).display());
                }
            }
        
            // Try one last approach - using a bind mount instead of overlay
            println!("\nTrying fallback approach with bind mount instead of overlay...");
            
            // First unmount any previous attempts
            let _ = run_with_privileges("umount", &[merged_dir]);
            let _ = run_with_privileges("umount", &[upper_dir]);
            
            // Remount tmpfs on upper_dir
            let tmpfs_status = run_with_privileges("mount", &["-t", "tmpfs", "tmpfs", upper_dir])
                .context("Failed mounting tmpfs on upper_dir")?;
            
            if !tmpfs_status.success() {
                println!("Failed to mount tmpfs for fallback approach.");
                
                // Clean up any directories we created
                remove_dir_all(merged_dir).ok();
                remove_dir_all(upper_dir).ok();
                
                println!("\nContinuing with limited functionality (no overlay mount).");
                
                // Instead of returning an error, continue with the rest of the program
                println!("Ephemeral overlay setup skipped.");
                println!("Cleanup complete. All ephemeral changes discarded.");
                return Ok(());
            }
            
            // Copy files from lower_dir to upper_dir
            println!("Copying files from {} to {}...", lower_dir, upper_dir);
            let copy_status = Command::new("cp")
                .args(&["-a", format!("{}/*", lower_dir).as_str(), upper_dir])
                .status();
                
            // Ignore errors from cp - the directory might be empty
            
            // Now bind mount upper_dir to merged_dir
            println!("Binding {} to {}...", upper_dir, merged_dir);
            let bind_status = run_with_privileges(
                "mount", 
                &["--bind", upper_dir, merged_dir]
            )
            .context("Failed to create bind mount")?;
            
            if !bind_status.success() {
                println!("Bind mount failed as well. All mount approaches have failed.");
                
                // Clean up any directories we created
                remove_dir_all(merged_dir).ok();
                remove_dir_all(upper_dir).ok();
                
                println!("\nContinuing with limited functionality (no mount).");
                
                // Instead of returning an error, continue with the rest of the program
                println!("Ephemeral overlay setup skipped.");
                println!("Cleanup complete. All ephemeral changes discarded.");
                return Ok(());
            }
            
            println!("Successfully created bind mount as fallback. Note that this is not a true overlay.");
            println!("Changes will be made directly to the upper directory.");
        }
    }
    println!("Ephemeral overlay set up at {}", merged_dir);

    // 6) Write a test file in the overlay
    let test_path = format!("{}/hello.txt", merged_dir);
    std::fs::OpenOptions::new()
        .create(true).write(true)
        .open(&test_path)?
        .write_all(b"Hello from ephemeral overlay!")?;
    println!("Wrote ephemeral file at {}", test_path);

    // Launch an interactive shell in the overlay environment
    launch_interactive_shell(merged_dir)?;

    // 7) Unmount for cleanup
    let umount_overlay = run_with_privileges("umount", &[merged_dir])
        .context("Failed to umount overlay_dir")?;
    
    if !umount_overlay.success() {
        println!("Warning: could not unmount overlay at {}", merged_dir);
    }

    let umount_upper = run_with_privileges("umount", &[upper_dir])
        .context("Failed to umount upper_dir")?;
    
    if !umount_upper.success() {
        println!("Warning: could not unmount tmpfs at {}", upper_dir);
    }

    remove_dir_all(merged_dir).ok();
    remove_dir_all(upper_dir).ok();
    // Optionally remove lower_dir if ephemeral too
    println!("Cleanup complete. All ephemeral changes discarded.");

    Ok(())
}