From 72db07051bea7fae8c1f4b46634b90838c3b335d Mon Sep 17 00:00:00 2001
From: iff <iff@ik.me>
Date: Mon, 31 Jul 2023 03:21:06 +0200
Subject: [PATCH] improv: small details

---
 rules/{sudo.toml => sudo-doas.toml} |  2 +-
 src/corrections.rs                  | 25 ++++++++++++-------------
 src/main.rs                         |  9 +++++++--
 src/shell.rs                        |  9 ++++++---
 src/style.rs                        |  1 -
 5 files changed, 26 insertions(+), 20 deletions(-)
 rename rules/{sudo.toml => sudo-doas.toml} (96%)

diff --git a/rules/sudo.toml b/rules/sudo-doas.toml
similarity index 96%
rename from rules/sudo.toml
rename to rules/sudo-doas.toml
index e00c99b..1bffe75 100644
--- a/rules/sudo.toml
+++ b/rules/sudo-doas.toml
@@ -1,4 +1,4 @@
-command = "sudo"
+command = "privilege"
 
 [[match_output]]
 pattern = [
diff --git a/src/corrections.rs b/src/corrections.rs
index 1c2fd93..37a4e0e 100644
--- a/src/corrections.rs
+++ b/src/corrections.rs
@@ -2,20 +2,20 @@ use std::collections::HashMap;
 
 use rule_parser::parse_rules;
 
-use crate::shell::command_output;
+use crate::shell::{command_output, PRIVILEGE_LIST};
 use crate::style::highlight_difference;
 
 pub fn correct_command(shell: &str, last_command: &str) -> Option<String> {
 	let command_output = command_output(shell, last_command);
 
 	let split_command = last_command.split_whitespace().collect::<Vec<&str>>();
-	let command = match split_command.first().expect("No command found.") {
-		&"sudo" => split_command.get(1).expect("No command found."),
-		_ => split_command.first().expect("No command found."),
+	let command = match PRIVILEGE_LIST.contains(&split_command[0]) {
+		true => split_command.get(1).expect("No command found."),
+		false => split_command.first().expect("No command found."),
 	};
 
-	if split_command[0] != "sudo" {
-		let suggest = match_pattern("sudo", &command_output);
+	if !PRIVILEGE_LIST.contains(command) {
+		let suggest = match_pattern("privilege", &command_output);
 		if let Some(suggest) = suggest {
 			let suggest = eval_suggest(&suggest, last_command);
 			return Some(suggest);
@@ -24,8 +24,8 @@ pub fn correct_command(shell: &str, last_command: &str) -> Option<String> {
 	let suggest = match_pattern(command, &command_output);
 	if let Some(suggest) = suggest {
 		let suggest = eval_suggest(&suggest, last_command);
-		if split_command[0] == "sudo" {
-			return Some(format!("sudo {}", suggest));
+		if PRIVILEGE_LIST.contains(command) {
+			return Some(format!("{} {}", split_command[0], suggest));
 		}
 		return Some(suggest);
 	}
@@ -92,12 +92,11 @@ 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();
 
-	let privilege = Vec::from(["doas", "sudo"]);
-
-	for p in privilege {
-		if command.starts_with(p){
+	for p in PRIVILEGE_LIST {
+		if command.starts_with(p) {
 			let command = command.replace(p, "");
-			std::process::Command::new(p)
+			println!("{} {}", p, command);
+			std::process::Command::new(p.trim())
 				.arg(shell)
 				.arg("-c")
 				.arg(command)
diff --git a/src/main.rs b/src/main.rs
index 632fa08..a134f5e 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -12,7 +12,7 @@ fn main() {
 	let corrected_command = corrections::correct_command(&shell, &last_command);
 
 	if let Some(mut corrected_command) = corrected_command {
-		if corrected_command.starts_with("sudo") {
+		if corrected_command.starts_with("sudo ") {
 			let privilege = get_privilege();
 			if let Some(privilege) = privilege {
 				if privilege != "sudo" {
@@ -22,6 +22,11 @@ fn main() {
 		}
 		corrections::confirm_correction(&shell, &corrected_command, &last_command);
 	} else {
-		println!("No correction found.");
+		println!(
+			"
+No correction found.
+
+If you think there should be a correction, please open an issue or send a pull request!"
+		);
 	}
 }
diff --git a/src/shell.rs b/src/shell.rs
index a3037cc..ac8568c 100644
--- a/src/shell.rs
+++ b/src/shell.rs
@@ -1,5 +1,7 @@
 use std::{collections::HashMap, fs::read_to_string, process::exit};
 
+pub const PRIVILEGE_LIST: [&str; 2] = ["doas ", "sudo "];
+
 pub fn find_shell() -> String {
 	std::env::var("SHELL")
 		.unwrap_or_else(|_| String::from("bash"))
@@ -17,7 +19,8 @@ pub fn find_last_command(shell: &str) -> String {
 		Err(_) => shell_default_history_file(shell),
 	};
 
-	let history = read_to_string(history_file).expect("Could not read history file.");
+	let history = read_to_string(history_file)
+		.expect("Could not read history file. Try setting the HISTFILE environment variable.");
 
 	match shell {
 		"bash" => history.lines().rev().nth(1).unwrap().to_string(),
@@ -73,8 +76,8 @@ fn shell_default_history_file(shell: &str) -> String {
 }
 
 pub fn get_privilege() -> Option<String> {
-	let list = vec!["doas", "sudo"];
-	for command in list {
+	for command in PRIVILEGE_LIST {
+		let command = command.trim();
 		if std::process::Command::new(command).output().is_ok() {
 			return Some(command.to_string());
 		}
diff --git a/src/style.rs b/src/style.rs
index eb8cbb8..a9b0568 100644
--- a/src/style.rs
+++ b/src/style.rs
@@ -1,4 +1,3 @@
-
 use colored::*;
 
 pub fn highlight_difference(corrected_command: &str, last_command: &str) -> String {