liblitho 0.2.0

cli tool to flash/clone the images to storage devices
Documentation
//! Linux privilege elevation via pkexec (polkit) or sudo.

use crate::platform::elevated_cli_args;
use crate::platform::traits::{ElevationCapability, PrivilegeOps};
use std::fs;
use std::io::{self, Write};
use std::process::{Command, Stdio};

use super::LinuxPlatform;

/// How we will elevate on Linux (chosen at call time).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Elevator {
    /// Desktop polkit agent → GUI password prompt via `pkexec`.
    Pkexec,
    /// Headless / no agent → terminal password prompt via `sudo`.
    Sudo,
}

impl PrivilegeOps for LinuxPlatform {
    fn elevation_method() -> &'static str {
        match select_elevator() {
            Some(Elevator::Pkexec) => "pkexec",
            Some(Elevator::Sudo) => "sudo",
            None => "none",
        }
    }

    fn is_elevated() -> bool {
        // SAFETY: geteuid has no preconditions and is always safe to call.
        unsafe { libc::geteuid() == 0 }
    }

    fn elevation_capability() -> ElevationCapability {
        if Self::is_elevated() {
            return ElevationCapability {
                ready: true,
                status_label: "elevated",
            };
        }

        match select_elevator() {
            Some(Elevator::Pkexec) => ElevationCapability {
                ready: true,
                status_label: "polkit ready",
            },
            Some(Elevator::Sudo) => ElevationCapability {
                ready: true,
                status_label: "sudo ready",
            },
            None => ElevationCapability {
                ready: false,
                status_label: "no elev. helper",
            },
        }
    }

    fn validate_elevation_prerequisites() -> Result<(), String> {
        if Self::is_elevated() {
            return Ok(());
        }
        match select_elevator() {
            Some(Elevator::Pkexec) | Some(Elevator::Sudo) => Ok(()),
            None => Err(
                "Cannot elevate: need either a polkit agent (pkexec) or sudo. \
                 On a headless system, install sudo and ensure this user can run sudo."
                    .into(),
            ),
        }
    }

    fn elevation_dialog_lines() -> [&'static str; 3] {
        match select_elevator() {
            Some(Elevator::Sudo) => [
                "Flash and clone operations require root access.",
                "This window will close; enter your sudo password in the terminal.",
                "After authentication, a new elevated session will start.",
            ],
            Some(Elevator::Pkexec) => [
                "Flash and clone operations require root access.",
                "You will be prompted for your password in a system dialog.",
                "This window will close and a new elevated session will start.",
            ],
            None => [
                "Flash and clone operations require root access.",
                "No elevation helper is available (need pkexec+agent or sudo).",
                "Install sudo for headless use, or start a polkit agent on desktop.",
            ],
        }
    }

    fn relaunch_elevated(mode: &str, device: &str, image: &str) -> Result<(), String> {
        use std::os::unix::process::CommandExt;

        let exe =
            std::env::current_exe().map_err(|e| format!("Could not resolve executable: {e}"))?;
        let args = elevated_cli_args(mode, device, image);

        if Self::is_elevated() {
            let mut cmd = Command::new(&exe);
            cmd.args(&args);
            preserve_terminal_env(&mut cmd);
            let err = cmd.exec();
            return Err(format!("Failed to exec elevated litho-tui: {err}"));
        }

        let elevator = select_elevator().ok_or_else(|| {
            "Cannot elevate: need either a polkit agent (pkexec) or sudo.".to_string()
        })?;

        let method = match elevator {
            Elevator::Pkexec => "pkexec",
            Elevator::Sudo => "sudo",
        };

        let mut cmd = match elevator {
            Elevator::Pkexec => {
                let mut cmd = Command::new("pkexec");
                cmd.arg(&exe).args(&args);
                cmd
            }
            Elevator::Sudo => {
                let _ = writeln!(
                    io::stderr(),
                    "\nlitho-tui: elevating with sudo — enter your password if prompted.\n"
                );
                let _ = io::stderr().flush();

                let mut cmd = Command::new("sudo");
                cmd.arg("--preserve-env=TERM,COLORTERM,LANG,LC_ALL,LC_CTYPE,COLORFGBG");
                cmd.arg(&exe).args(&args);
                cmd
            }
        };

        preserve_terminal_env(&mut cmd);

        let status = cmd
            .status()
            .map_err(|e| format!("Failed to run {method}: {e}"))?;

        if status.success() {
            std::process::exit(0);
        }

        let code = {
            use std::os::unix::process::ExitStatusExt;
            status.code().map(|c| c.to_string()).unwrap_or_else(|| {
                status
                    .signal()
                    .map(|s| format!("signal {s}"))
                    .unwrap_or_else(|| "unknown".into())
            })
        };

        Err(match elevator {
            Elevator::Pkexec => format!(
                "Elevation cancelled or failed (pkexec exited: {code}). \
                 Is a polkit agent running?"
            ),
            Elevator::Sudo => format!(
                "Elevation cancelled or failed (sudo exited: {code}). \
                 Check your password, or that this user may run sudo."
            ),
        })
    }

    fn current_euid_or_0() -> u32 {
        // SAFETY: geteuid has no preconditions and is always safe to call.
        unsafe { libc::geteuid() }
    }
}

fn select_elevator() -> Option<Elevator> {
    let pkexec = pkexec_on_path();
    let sudo = sudo_on_path();
    let agent_running = find_running_polkit_auth_agent().is_some();
    let agent_installed = find_polkit_auth_agent_binary().is_some();

    if pkexec && agent_running {
        return Some(Elevator::Pkexec);
    }
    if sudo {
        return Some(Elevator::Sudo);
    }
    if pkexec && agent_installed {
        return Some(Elevator::Pkexec);
    }
    if pkexec {
        return Some(Elevator::Pkexec);
    }
    None
}

pub(crate) fn pkexec_on_path() -> bool {
    which_ok("pkexec")
}

fn sudo_on_path() -> bool {
    which_ok("sudo")
}

fn which_ok(bin: &str) -> bool {
    Command::new("which")
        .arg(bin)
        .output()
        .map(|o| o.status.success())
        .unwrap_or(false)
}

fn current_username() -> String {
    std::env::var("USER")
        .or_else(|_| std::env::var("LOGNAME"))
        .unwrap_or_else(|_| "root".to_string())
}

fn extract_executable_path_from_ps(line: &str) -> Option<String> {
    for token in line.split_whitespace() {
        if token.starts_with('/') && token.contains("polkit") {
            return Some(token.to_string());
        }
    }
    None
}

fn find_running_polkit_auth_agent() -> Option<String> {
    let user = current_username();

    let output = Command::new("ps")
        .args(["-u", &user, "-o", "pid,comm,args", "--no-headers"])
        .output()
        .ok()?;

    let stdout = String::from_utf8_lossy(&output.stdout);
    for line in stdout.lines() {
        let lower = line.to_lowercase();
        if lower.contains("polkit")
            && (lower.contains("agent") || lower.contains("-authentication-agent"))
        {
            if let Some(path) = extract_executable_path_from_ps(line) {
                if fs::metadata(&path).map(|m| m.is_file()).unwrap_or(false) {
                    return Some(path);
                }
            }
            return Some(line.trim().to_string());
        }
    }
    None
}

fn find_polkit_auth_agent_binary() -> Option<String> {
    const CANDIDATES: &[&str] = &[
        "/usr/libexec/polkit-gnome-authentication-agent-1",
        "/usr/lib/polkit-gnome/polkit-gnome-authentication-agent-1",
        "/usr/libexec/polkit-kde-authentication-agent-1",
        "/usr/lib/polkit-kde-authentication-agent-1",
        "/usr/lib/x86_64-linux-gnu/libexec/polkit-kde-authentication-agent-1",
        "/usr/libexec/xfce-polkit",
        "/usr/libexec/polkit-mate-authentication-agent-1",
        "/usr/bin/lxpolkit",
        "/usr/libexec/cinnamon-polkit",
    ];

    for path in CANDIDATES {
        if fs::metadata(path).map(|m| m.is_file()).unwrap_or(false) {
            return Some(path.to_string());
        }
    }
    None
}

fn preserve_terminal_env(cmd: &mut Command) {
    cmd.stdin(Stdio::inherit())
        .stdout(Stdio::inherit())
        .stderr(Stdio::inherit());

    for var in ["TERM", "COLORTERM", "LANG", "LC_ALL", "LC_CTYPE", "COLORFGBG"] {
        if let Ok(value) = std::env::var(var) {
            cmd.env(var, value);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::platform::traits::PrivilegeOps;

    #[test]
    fn select_elevator_never_panics() {
        let _ = select_elevator();
    }

    #[test]
    fn which_helpers_are_consistent() {
        if matches!(select_elevator(), Some(Elevator::Sudo)) {
            assert!(sudo_on_path());
        }
        if matches!(select_elevator(), Some(Elevator::Pkexec)) {
            assert!(pkexec_on_path());
        }
    }

    #[test]
    fn find_polkit_agent_running_or_binary_is_optional() {
        let _ = find_running_polkit_auth_agent();
        let _ = find_polkit_auth_agent_binary();
    }

    #[test]
    fn elevated_process_reports_capability_ready() {
        if LinuxPlatform::is_elevated() {
            assert!(LinuxPlatform::elevation_capability().ready);
        }
    }

    #[test]
    fn polkit_ready_implies_pkexec_on_path() {
        let cap = LinuxPlatform::elevation_capability();
        if cap.ready && cap.status_label == "polkit ready" {
            assert!(pkexec_on_path());
            assert_eq!(LinuxPlatform::elevation_method(), "pkexec");
        }
    }

    #[test]
    fn sudo_ready_implies_sudo_method() {
        let cap = LinuxPlatform::elevation_capability();
        if cap.ready && cap.status_label == "sudo ready" {
            assert_eq!(LinuxPlatform::elevation_method(), "sudo");
        }
    }
}