diff --git a/src/args.rs b/src/args.rs index fd7d9ab..529193c 100644 --- a/src/args.rs +++ b/src/args.rs @@ -1,9 +1,10 @@ use crate::shell::initialization; -pub fn handle_args() { +// returns true if should exit +pub fn handle_args() -> bool { let args = std::env::args().collect::>(); if args.len() <= 1 { - return; + return false; } let mut auto_aliasing = String::new(); let mut shell = String::new(); @@ -13,9 +14,11 @@ pub fn handle_args() { match args[index].as_str() { "-h" | "--help" => { print_help(); + return true; } "-v" | "--version" => { print_version(); + return true; } "-a" | "--alias" => { if args.len() > index + 1 { @@ -43,12 +46,13 @@ pub fn handle_args() { if shell.is_empty() { eprintln!("{}", t!("no-shell")); - std::process::exit(1); + return true; } let binary_path = &args[0]; initialization(&shell, binary_path, &auto_aliasing, cnf); + return true; } fn print_help() { @@ -63,7 +67,6 @@ fn print_help() { auto_example_fish = "pay-respects fish --alias | source", ) ); - std::process::exit(0); } fn print_version() { @@ -84,5 +87,4 @@ fn print_version() { { println!(" - libcurl"); } - std::process::exit(0); } diff --git a/src/main.rs b/src/main.rs index 41bb573..68a8e60 100644 --- a/src/main.rs +++ b/src/main.rs @@ -37,9 +37,28 @@ mod requests; extern crate rust_i18n; i18n!("i18n", fallback = "en", minify_key = true); -fn main() { +fn main() -> Result<(), std::io::Error>{ colored::control::set_override(true); - // let locale = std::env::var("LANG").unwrap_or("en_US".to_string()); + let mut data = { + let init = init(); + if init.is_none() { + return Ok(()); + } else { + init.unwrap() + } + }; + + data.expand_command(); + use shell::Mode; + match data.mode { + Mode::Suggestion => modes::suggestion(&mut data), + Mode::CNF => modes::cnf(&mut data), + } + + Ok(()) +} + +fn init() -> Option { let locale = { let sys_locale = get_locale().unwrap_or("en-US".to_string()); if sys_locale.len() < 2 { @@ -50,6 +69,11 @@ fn main() { }; rust_i18n::set_locale(&locale[0..2]); + let exit = args::handle_args(); + if exit { + return None; + } + #[cfg(feature = "request-ai")] { if std::env::var("_PR_AI_LOCALE").is_err() { @@ -57,16 +81,5 @@ fn main() { } } - args::handle_args(); - - let mode = match std::env::var("_PR_MODE") { - Ok(mode) => mode, - Err(_) => "suggestion".to_string(), - }; - - match mode.as_str() { - "suggestion" => modes::suggestion(), - "cnf" => modes::cnf(), - _ => {} - } + Some(shell::Data::init()) } diff --git a/src/modes.rs b/src/modes.rs index 90e354d..a49fb4f 100644 --- a/src/modes.rs +++ b/src/modes.rs @@ -1,32 +1,18 @@ -use crate::shell::{command_output, get_shell, PRIVILEGE_LIST}; +use crate::shell::Data; use crate::style::highlight_difference; -use crate::suggestions::{best_match_path, split_command}; +use crate::suggestions::{best_match_path, suggest_command}; use crate::system; use crate::{shell, suggestions}; use colored::Colorize; use inquire::*; -pub fn suggestion() { - let shell = get_shell(); - let mut last_command = shell::last_command(&shell).trim().to_string(); - last_command = shell::expand_alias(&shell, &last_command); - let mut error_msg = { - let error_msg = std::env::var("_PR_ERROR_MSG"); - if let Ok(error_msg) = error_msg { - error_msg - } else { - command_output(&shell, &last_command) - } - }; - - error_msg = error_msg - .split_whitespace() - .collect::>() - .join(" "); +pub fn suggestion(data: &mut Data) { + let shell = data.shell.clone(); + let last_command = data.command.clone(); loop { let suggestion = { - let command = suggestions::suggest_command(&shell, &last_command, &error_msg); + let command = suggest_command(&data); if command.is_none() { break; }; @@ -35,6 +21,8 @@ pub fn suggestion() { shell::shell_syntax(&shell, &mut command); command }; + data.update_suggest(&suggestion); + data.expand_suggest(); let highlighted_suggestion = { let difference = highlight_difference(&shell, &suggestion, &last_command); @@ -45,16 +33,13 @@ pub fn suggestion() { }; let execution = - suggestions::confirm_suggestion(&shell, &suggestion, &highlighted_suggestion); + suggestions::confirm_suggestion(&data, &highlighted_suggestion); if execution.is_ok() { return; } else { - last_command = suggestion; - error_msg = execution.err().unwrap(); - error_msg = error_msg - .split_whitespace() - .collect::>() - .join(" "); + data.update_command(&suggestion); + let msg = Some(execution.err().unwrap().split_whitespace().collect::>().join(" ")); + data.update_error(msg); let retry_message = format!("{}...", t!("retry")); @@ -69,33 +54,24 @@ pub fn suggestion() { ); } -pub fn cnf() { - let shell = get_shell(); - let mut last_command = shell::last_command(&shell).trim().to_string(); - last_command = shell::expand_alias(&shell, &last_command); +pub fn cnf(data: &mut Data) { + let shell = data.shell.clone(); + let last_command = data.command.clone(); + let mut split_command = data.split.clone(); - let mut split_command = split_command(&last_command); - let executable = match PRIVILEGE_LIST.contains(&split_command[0].as_str()) { - true => split_command.get(1).expect(&t!("no-command")).as_str(), - false => split_command.first().expect(&t!("no-command")).as_str(), - }; + let executable = split_command[0].as_str(); let best_match = best_match_path(executable); if best_match.is_some() { let best_match = best_match.unwrap(); - match PRIVILEGE_LIST.contains(&split_command[0].as_str()) { - true => { - split_command[1] = best_match; - } - false => { - split_command[0] = best_match; - } - } + split_command[0] = best_match; let suggestion = split_command.join(" "); + data.update_suggest(&suggestion); + data.expand_suggest(); let highlighted_suggestion = highlight_difference(&shell, &suggestion, &last_command).unwrap(); - let _ = suggestions::confirm_suggestion(&shell, &suggestion, &highlighted_suggestion); + let _ = suggestions::confirm_suggestion(&data, &highlighted_suggestion); } else { let package_manager = match system::get_package_manager(&shell) { Some(package_manager) => package_manager, @@ -124,7 +100,7 @@ pub fn cnf() { // retry after installing package if system::install_package(&shell, &package_manager, &package) { - let _ = suggestions::confirm_suggestion(&shell, &last_command, &last_command); + let _ = suggestions::confirm_suggestion(&data, &last_command); } } } diff --git a/src/shell.rs b/src/shell.rs index 35242e2..ebd1939 100644 --- a/src/shell.rs +++ b/src/shell.rs @@ -3,9 +3,137 @@ use std::process::exit; use std::sync::mpsc::channel; use std::thread; use std::time::Duration; +use std::collections::HashMap; + +use regex_lite::Regex; pub const PRIVILEGE_LIST: [&str; 2] = ["sudo", "doas"]; +pub enum Mode { + Suggestion, + CNF, +} + +pub struct Data { + pub shell: String, + pub command: String, + pub suggest: Option, + pub split: Vec, + pub alias: Option>, + pub privilege: Option, + pub error: String, + pub mode: Mode, +} + +pub struct Init { + pub shell: String, + pub binary_path: String, + pub auto_alias: String, + pub cnf: bool, +} + +impl Data { + pub fn init() -> Data { + let shell = get_shell(); + let command = last_command(&shell).trim().to_string(); + let alias = alias_map(&shell); + let mode = run_mode(); + + let mut init = Data { + shell, + command, + suggest: None, + alias, + split: vec![], + privilege: None, + error: "".to_string(), + mode, + }; + + init.split(); + init.update_error(None); + init + } + + pub fn expand_command(&mut self) { + if self.alias.is_none() { + return; + } + let alias = self.alias.as_ref().unwrap(); + if let Some(command) = expand_alias_multiline(alias, &self.command) { + self.update_command(&command); + } + } + + pub fn expand_suggest(&mut self) { + if self.alias.is_none() { + return; + } + let alias = self.alias.as_ref().unwrap(); + if let Some(suggest) = expand_alias_multiline(alias, &self.suggest.as_ref().unwrap()) { + self.update_suggest(&suggest); + } + } + + pub fn split(&mut self) { + let mut split = split_command(&self.command); + if PRIVILEGE_LIST.contains(&split[0].as_str()) { + self.command = self.command.replacen(&split[0], "", 1); + self.privilege = Some(split.remove(0)) + } + self.split = split; + } + + pub fn update_error(&mut self, error: Option) { + if let Some(error) = error { + self.error = error; + } else { + self.error = get_error(&self.shell, &self.command); + } + } + + pub fn update_command(&mut self, command: &str) { + self.command = command.to_string(); + self.split(); + } + + pub fn update_suggest(&mut self, suggest: &str) { + let split = split_command(&suggest); + if PRIVILEGE_LIST.contains(&split[0].as_str()) { + self.suggest = Some(suggest.replacen(&split[0], "", 1)); + self.privilege = Some(split[0].clone()) + } else { + self.suggest = Some(suggest.to_string()); + }; + } +} + +pub fn split_command(command: &str) -> Vec { + if cfg!(debug_assertions) { + eprintln!("command: {command}") + } + // this regex splits the command separated by spaces, except when the space + // is escaped by a backslash or surrounded by quotes + let regex = r#"([^\s"'\\]+|"(?:\\.|[^"\\])*"|'(?:\\.|[^'\\])*'|\\ )+|\\|\n"#; + let regex = Regex::new(regex).unwrap(); + let split_command = regex + .find_iter(command) + .map(|cap| cap.as_str().to_owned()) + .collect::>(); + split_command +} + +pub fn get_error(shell: &str, command: &str) -> String { + let error_msg = std::env::var("_PR_ERROR_MSG"); + let error = if let Ok(error_msg) = error_msg { + std::env::remove_var("_PR_ERROR_MSG"); + error_msg + } else { + command_output(shell, command) + }; + error.split_whitespace().collect::>().join(" ") +} + pub fn command_output(shell: &str, command: &str) -> String { let (sender, receiver) = channel(); @@ -65,85 +193,93 @@ pub fn last_command(shell: &str) -> String { } } -pub fn expand_alias(shell: &str, full_command: &str) -> String { - let alias_env = std::env::var("_PR_ALIAS"); - if alias_env.is_err() { - return full_command.to_string(); +pub fn run_mode() -> Mode { + match std::env::var("_PR_MODE") { + Ok(mode) => { + match mode.as_str() { + "suggestion" => Mode::Suggestion, + "cnf" => Mode::CNF, + _ => Mode::Suggestion, + } + } + Err(_) => Mode::Suggestion, } - let alias = alias_env.unwrap(); - if alias.is_empty() { - return full_command.to_string(); +} + +pub fn alias_map(shell: &str) -> Option> { + let env = std::env::var("_PR_ALIAS"); + if env.is_err() { + return None; + } + let env = env.unwrap(); + if env.is_empty() { + return None; } - let split_command = full_command.split_whitespace().collect::>(); - let (command, pure_command) = if PRIVILEGE_LIST.contains(&split_command[0]) && split_command.len() > 1 { - (split_command[1], Some(split_command[1..].join(" "))) - } else { - (split_command[0], None) - }; - - let mut expanded_command = Option::None; - + let mut alias_map = HashMap::new(); match shell { "bash" => { - for line in alias.lines() { - if line.starts_with(format!("alias {}=", command).as_str()) { - let alias = line.replace(format!("alias {}='", command).as_str(), ""); - let alias = alias.trim_end_matches('\'').trim_start_matches('\''); - - expanded_command = Some(alias.to_string()); - } + for line in env.lines() { + let alias = line.replace("alias ", ""); + let (alias, command) = alias.split_once('=').unwrap(); + let command = command.trim().trim_matches('\''); + alias_map.insert(alias.to_string(), command.to_string()); } } "zsh" => { - for line in alias.lines() { - if line.starts_with(format!("{}=", command).as_str()) { - let alias = line.replace(format!("{}=", command).as_str(), ""); - let alias = alias.trim_start_matches('\'').trim_end_matches('\''); - - expanded_command = Some(alias.to_string()); - } + for line in env.lines() { + let (alias, command) = line.split_once('=').unwrap(); + let command = command.trim().trim_matches('\''); + alias_map.insert(alias.to_string(), command.to_string()); } } "fish" => { - for line in alias.lines() { - if line.starts_with(format!("alias {} ", command).as_str()) { - let alias = line.replace(format!("alias {} ", command).as_str(), ""); - let alias = alias.trim_start_matches('\'').trim_end_matches('\''); - - expanded_command = Some(alias.to_string()); - } + for line in env.lines() { + let alias = line.replace("alias ", ""); + let (alias, command) = alias.split_once(' ').unwrap(); + let command = command.trim().trim_matches('\''); + alias_map.insert(alias.to_string(), command.to_string()); } } _ => { - eprintln!("Unsupported shell: {}", shell); - exit(1); - } - }; - - if expanded_command.is_none() { - return full_command.to_string(); - }; - - let expanded_command = expanded_command.unwrap(); - - if pure_command.is_some() { - let pure_command = pure_command.unwrap(); - if pure_command.starts_with(&expanded_command) { - return full_command.to_string(); + unreachable!("Unsupported shell: {}", shell); } } - - full_command.replacen(command, &expanded_command, 1) + Some(alias_map) } -pub fn expand_alias_multiline(shell: &str, full_command: &str) -> String { - let lines = full_command.lines().collect::>(); - let mut expanded = String::new(); - for line in lines { - expanded = format!("{}\n{}", expanded, expand_alias(shell, line)); +pub fn expand_alias(map: &HashMap, command: &str) -> Option { + #[cfg(debug_assertions)] + eprintln!("expand_alias: command: {}", command); + let command = if let Some(split) = command.split_once(' ') { + split.0 + } else { + command + }; + if let Some(expand) = map.get(command) { + Some(command.replacen(&command, expand, 1)) + } else { + None + } +} + +pub fn expand_alias_multiline(map: &HashMap, command: &str) -> Option { + let lines = command.lines().collect::>(); + let mut expanded = String::new(); + let mut expansion = false; + for line in lines { + if let Some(expand) = expand_alias(&map, line){ + expanded = format!("{}\n{}", expanded, expand); + expansion = true; + } else { + expanded = format!("{}\n{}", expanded, line); + } + } + if expansion { + Some(expanded) + } else { + None } - expanded } pub fn initialization(shell: &str, binary_path: &str, auto_alias: &str, cnf: bool) { @@ -173,7 +309,7 @@ pub fn initialization(shell: &str, binary_path: &str, auto_alias: &str, cnf: boo } _ => { println!("Unknown shell: {}", shell); - std::process::exit(1); + return; } } @@ -194,7 +330,7 @@ def --env {} [] {{ pr_alias, last_command, binary_path ); println!("{}", init); - std::process::exit(0); + return; } let mut init = match shell { @@ -233,12 +369,12 @@ def --env {} [] {{ ), _ => { println!("Unsupported shell: {}", shell); - exit(1); + return; } }; if auto_alias.is_empty() { println!("{}", init); - std::process::exit(0); + return; } match shell { @@ -264,7 +400,7 @@ end } _ => { println!("Unsupported shell: {}", shell); - exit(1); + return; } } @@ -332,14 +468,12 @@ $ExecutionContext.InvokeCommand.CommandNotFoundAction = } _ => { println!("Unsupported shell: {}", shell); - exit(1); + return; } } } println!("{}", init); - - std::process::exit(0); } pub fn get_shell() -> String { diff --git a/src/suggestions.rs b/src/suggestions.rs index 2c5a889..0be11a1 100644 --- a/src/suggestions.rs +++ b/src/suggestions.rs @@ -7,52 +7,39 @@ use regex_lite::Regex; use crate::files::{get_best_match_file, get_path_files}; use crate::rules::match_pattern; -use crate::shell::{expand_alias_multiline, shell_evaluated_commands, PRIVILEGE_LIST}; +use crate::shell::{Data, shell_evaluated_commands}; -pub fn suggest_command(shell: &str, last_command: &str, error_msg: &str) -> Option { - let split_command = split_command(last_command); - let executable = match PRIVILEGE_LIST.contains(&split_command[0].as_str()) { - true => split_command.get(1).expect(&t!("no-command")).as_str(), - false => split_command.first().expect(&t!("no-command")).as_str(), - }; +pub fn suggest_command(data: &Data) -> Option { + let shell = &data.shell; + let command = &data.command; + let split_command = &data.split; + let executable = data.split[0].as_str(); + let error = &data.error; + let privilege = &data.privilege; - if !PRIVILEGE_LIST.contains(&executable) { - let suggest = match_pattern("_PR_privilege", last_command, error_msg, shell); + if privilege.is_none() { + let suggest = match_pattern("_PR_privilege", command, error, shell); if suggest.is_some() { return suggest; } } - let last_command = match PRIVILEGE_LIST.contains(&split_command[0].as_str()) { - true => &last_command[split_command[0].len() + 1..], - false => last_command, - }; - - let suggest = match_pattern(executable, last_command, error_msg, shell); - if let Some(suggest) = suggest { - if PRIVILEGE_LIST.contains(&split_command[0].as_str()) { - return Some(format!("{} {}", split_command[0], suggest)); - } - return Some(suggest); + let suggest = match_pattern(executable, command, error, shell); + if suggest.is_some() { + return suggest; } - let suggest = match_pattern("_PR_general", last_command, error_msg, shell); - if let Some(suggest) = suggest { - if PRIVILEGE_LIST.contains(&split_command[0].as_str()) { - return Some(format!("{} {}", split_command[0], suggest)); - } - return Some(suggest); + let suggest = match_pattern("_PR_general", command, error, shell); + if suggest.is_some() { + return suggest; } #[cfg(feature = "runtime-rules")] { use crate::runtime_rules::runtime_match; - let suggest = runtime_match(executable, last_command, error_msg, shell); - if let Some(suggest) = suggest { - if PRIVILEGE_LIST.contains(&split_command[0].as_str()) { - return Some(format!("{} {}", split_command[0], suggest)); - } - return Some(suggest); + let suggest = runtime_match(executable, command, error, shell); + if suggest.is_some() { + return suggest; } } @@ -63,22 +50,17 @@ pub fn suggest_command(shell: &str, last_command: &str, error_msg: &str) -> Opti // skip for commands with no arguments, // very likely to be an error showing the usage - if PRIVILEGE_LIST.contains(&split_command[0].as_str()) && split_command.len() > 2 - || !PRIVILEGE_LIST.contains(&split_command[0].as_str()) && split_command.len() > 1 + if privilege.is_some() && split_command.len() > 2 || + privilege.is_none() && split_command.len() > 1 { - let suggest = ai_suggestion(last_command, error_msg); + let suggest = ai_suggestion(command, error); if let Some(suggest) = suggest { let warn = format!("{}:", t!("ai-suggestion")).bold().blue(); let note = fill(&suggest.note, termwidth()); eprintln!("{}\n{}\n", warn, note); let command = suggest.command; - if command != "None" { - if PRIVILEGE_LIST.contains(&split_command[0].as_str()) { - return Some(format!("{} {}", split_command[0], command)); - } - return Some(command); - } + return Some(command); } } } @@ -262,56 +244,17 @@ pub fn compare_string(a: &str, b: &str) -> usize { matrix[a.chars().count()][b.chars().count()] } -pub fn confirm_suggestion(shell: &str, command: &str, highlighted: &str) -> Result<(), String> { +pub fn confirm_suggestion(data: &Data, highlighted: &str) -> Result<(), String> { eprintln!("{}\n", highlighted); let confirm = format!("[{}]", t!("confirm-yes")).green(); eprintln!("{}: {} {}", t!("confirm"), confirm, "[Ctrl+C]".red()); std::io::stdin().read_line(&mut String::new()).unwrap(); - for p in PRIVILEGE_LIST { - let _p = p.to_owned() + " "; - if !command.starts_with(&_p) { - continue; - } + let shell = &data.shell; + let command = &data.suggest.clone().unwrap(); - let command = { - let mut command = command.replacen(&_p, "", 1); - if command != " " { - command = expand_alias_multiline(shell, &command); - } - command - }; - - let now = Instant::now(); - let process = run_suggestion_p(shell, p, &command); - - if process.success() { - let cd = shell_evaluated_commands(shell, &command); - if let Some(cd) = cd { - println!("{}", cd); - } - } else { - if now.elapsed() > Duration::from_secs(3) { - exit(1); - } - let process = std::process::Command::new(p) - .arg(shell) - .arg("-c") - .arg(command) - .env("LC_ALL", "C") - .output() - .expect("failed to execute process"); - let error_msg = match process.stderr.is_empty() { - true => String::from_utf8_lossy(&process.stdout), - false => String::from_utf8_lossy(&process.stderr), - }; - return Err(error_msg.to_string()); - } - } - - let command = expand_alias_multiline(shell, command); let now = Instant::now(); - let process = run_suggestion(shell, &command); + let process = run_suggestion(data, &command); if process.success() { let cd = shell_evaluated_commands(shell, &command); @@ -337,28 +280,32 @@ pub fn confirm_suggestion(shell: &str, command: &str, highlighted: &str) -> Resu } } -fn run_suggestion_p(shell: &str, p: &str, command: &str) -> std::process::ExitStatus { - std::process::Command::new(p) - .arg(shell) - .arg("-c") - .arg(command) - .stdout(stderr()) - .stderr(Stdio::inherit()) - .spawn() - .expect("failed to execute process") - .wait() - .unwrap() -} - -fn run_suggestion(shell: &str, command: &str) -> std::process::ExitStatus { - std::process::Command::new(shell) - .arg("-c") - .arg(command) - // .stdout(Stdio::inherit()) - .stdout(stderr()) - .stderr(Stdio::inherit()) - .spawn() - .expect("failed to execute process") - .wait() - .unwrap() +fn run_suggestion(data: &Data, command: &str) -> std::process::ExitStatus { + let shell = &data.shell; + let privilege = &data.privilege; + match privilege { + Some(sudo) => { + std::process::Command::new(sudo) + .arg(shell) + .arg("-c") + .arg(command) + .stdout(stderr()) + .stderr(Stdio::inherit()) + .spawn() + .expect("failed to execute process") + .wait() + .unwrap() + } + None => { + std::process::Command::new(shell) + .arg("-c") + .arg(command) + .stdout(stderr()) + .stderr(Stdio::inherit()) + .spawn() + .expect("failed to execute process") + .wait() + .unwrap() + } + } }