diff --git a/src/corrections.rs b/src/corrections.rs index 2d789a4..a25b1bd 100644 --- a/src/corrections.rs +++ b/src/corrections.rs @@ -92,22 +92,28 @@ pub fn confirm_correction(shell: &str, command: &str, last_command: &str) { println!("Press enter to execute the corrected command. Or press Ctrl+C to exit."); std::io::stdin().read_line(&mut String::new()).unwrap(); - if command.starts_with("sudo") { - std::process::Command::new("sudo") - .arg(shell) - .arg("-c") - .arg(command) - .spawn() - .expect("failed to execute process") - .wait() - .expect("failed to wait on process"); - } else { - std::process::Command::new(shell) - .arg("-c") - .arg(command) - .spawn() - .expect("failed to execute process") - .wait() - .expect("failed to wait on process"); + let privilege = Vec::from(["sudo", "doas"]); + + for p in privilege { + if command.starts_with(p){ + let command = command.replace(p, ""); + std::process::Command::new(p) + .arg(shell) + .arg("-c") + .arg(command) + .spawn() + .expect("failed to execute process") + .wait() + .expect("failed to wait on process"); + return; + } } + + std::process::Command::new(shell) + .arg("-c") + .arg(command) + .spawn() + .expect("failed to execute process") + .wait() + .expect("failed to wait on process"); } diff --git a/src/main.rs b/src/main.rs index 5764654..632fa08 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,13 +2,24 @@ mod corrections; mod shell; mod style; +use shell::get_privilege; + fn main() { std::env::set_var("LC_ALL", "C"); let shell = shell::find_shell(); let last_command = shell::find_last_command(&shell); let corrected_command = corrections::correct_command(&shell, &last_command); - if let Some(corrected_command) = corrected_command { + + if let Some(mut corrected_command) = corrected_command { + if corrected_command.starts_with("sudo") { + let privilege = get_privilege(); + if let Some(privilege) = privilege { + if privilege != "sudo" { + corrected_command = corrected_command.replacen("sudo", &privilege, 1); + } + } + } corrections::confirm_correction(&shell, &corrected_command, &last_command); } else { println!("No correction found."); diff --git a/src/shell.rs b/src/shell.rs index f70b024..4b5b448 100644 --- a/src/shell.rs +++ b/src/shell.rs @@ -71,3 +71,13 @@ fn shell_default_history_file(shell: &str) -> String { let file = shell_file_map.get(shell).expect("Unsupported shell."); format!("{}/{}", std::env::var("HOME").unwrap(), file) } + +pub fn get_privilege() -> Option { + let list = vec!["sudo", "doas"]; + for command in list { + if std::process::Command::new(command).output().is_ok() { + return Some(command.to_string()); + } + } + None +} diff --git a/src/style.rs b/src/style.rs index a9b0568..eb8cbb8 100644 --- a/src/style.rs +++ b/src/style.rs @@ -1,3 +1,4 @@ + use colored::*; pub fn highlight_difference(corrected_command: &str, last_command: &str) -> String {