From 0fba20173902e43fc8e7207abe8c58b59438625b Mon Sep 17 00:00:00 2001 From: Gal Schlezinger Date: Sun, 30 Jan 2022 14:00:32 +0200 Subject: [PATCH] Revert from using sysinfo, but only on Unix machines (#625) * Add the `ps` code back * some quality improvements * cargo clippy --fix * fix * correct path --- src/shell/infer/mod.rs | 45 +++++--------- src/shell/infer/unix.rs | 121 +++++++++++++++++++++++++++++++++++++ src/shell/infer/windows.rs | 34 +++++++++++ 3 files changed, 171 insertions(+), 29 deletions(-) create mode 100644 src/shell/infer/unix.rs create mode 100644 src/shell/infer/windows.rs diff --git a/src/shell/infer/mod.rs b/src/shell/infer/mod.rs index 157f042..e0c22c6 100644 --- a/src/shell/infer/mod.rs +++ b/src/shell/infer/mod.rs @@ -1,34 +1,21 @@ -use super::{Bash, Fish, PowerShell, Shell, WindowsCmd, Zsh}; -use log::debug; -use std::ffi::OsStr; -use sysinfo::{ProcessExt, System, SystemExt}; +mod unix; -pub fn infer_shell() -> Option> { - let mut system = System::new(); - let mut current_pid = sysinfo::get_current_pid().ok(); +mod windows; - while let Some(pid) = current_pid { - system.refresh_process(pid); - if let Some(process) = system.process(pid) { - current_pid = process.parent(); - let process_name = process - .exe() - .file_stem() - .and_then(OsStr::to_str) - .map(str::to_lowercase); - let sliced = process_name.as_ref().map(|x| &x[..]); - match sliced { - Some("sh" | "bash") => return Some(Box::from(Bash)), - Some("zsh") => return Some(Box::from(Zsh)), - Some("fish") => return Some(Box::from(Fish)), - Some("pwsh" | "powershell") => return Some(Box::from(PowerShell)), - Some("cmd") => return Some(Box::from(WindowsCmd)), - cmd_name => debug!("binary is not a supported shell: {:?}", cmd_name), - }; - } else { - current_pid = None; - } - } +#[cfg(unix)] +pub use self::unix::infer_shell; +#[cfg(not(unix))] +pub use self::windows::infer_shell; +pub(self) fn shell_from_string(shell: &str) -> Option> { + use super::{Bash, Fish, PowerShell, WindowsCmd, Zsh}; + match shell { + "sh" | "bash" => return Some(Box::from(Bash)), + "zsh" => return Some(Box::from(Zsh)), + "fish" => return Some(Box::from(Fish)), + "pwsh" | "powershell" => return Some(Box::from(PowerShell)), + "cmd" => return Some(Box::from(WindowsCmd)), + cmd_name => log::debug!("binary is not a supported shell: {:?}", cmd_name), + }; None } diff --git a/src/shell/infer/unix.rs b/src/shell/infer/unix.rs new file mode 100644 index 0000000..9e24dcf --- /dev/null +++ b/src/shell/infer/unix.rs @@ -0,0 +1,121 @@ +#![cfg(unix)] + +use crate::shell::Shell; +use log::debug; +use std::io::{Error, ErrorKind}; +use thiserror::Error; + +#[derive(Debug)] +struct ProcessInfo { + parent_pid: Option, + command: String, +} + +const MAX_ITERATIONS: u8 = 10; + +pub fn infer_shell() -> Option> { + let mut pid = Some(std::process::id()); + let mut visited = 0; + + while let Some(current_pid) = pid { + if visited > MAX_ITERATIONS { + return None; + } + + let process_info = get_process_info(current_pid) + .map_err(|err| { + debug!("{}", err); + err + }) + .ok()?; + let binary = process_info + .command + .trim_start_matches('-') + .split('/') + .last()?; + + if let Some(shell) = super::shell_from_string(binary) { + return Some(shell); + } + + pid = process_info.parent_pid; + visited += 1; + } + + None +} + +fn get_process_info(pid: u32) -> Result { + use std::io::{BufRead, BufReader}; + use std::process::Command; + + let buffer = Command::new("ps") + .arg("-o") + .arg("ppid,comm") + .arg(pid.to_string()) + .stdout(std::process::Stdio::piped()) + .spawn()? + .stdout + .ok_or_else(|| Error::from(ErrorKind::UnexpectedEof))?; + + let mut lines = BufReader::new(buffer).lines(); + + // skip header line + lines + .next() + .ok_or_else(|| Error::from(ErrorKind::UnexpectedEof))??; + + let line = lines + .next() + .ok_or_else(|| Error::from(ErrorKind::NotFound))??; + + let mut parts = line.trim().split_whitespace(); + let ppid = parts.next().ok_or_else(|| ProcessInfoError::Parse { + expectation: "Can't read the ppid from ps, should be the first item in the table", + got: line.to_string(), + })?; + let command = parts.next().ok_or_else(|| ProcessInfoError::Parse { + expectation: "Can't read the command from ps, should be the second item in the table", + got: line.to_string(), + })?; + + Ok(ProcessInfo { + parent_pid: ppid.parse().ok(), + command: command.into(), + }) +} + +#[derive(Debug, Error)] +enum ProcessInfoError { + #[error("Can't read process info: {source}")] + Io { + #[source] + #[from] + source: std::io::Error, + }, + #[error("Can't parse process info output. {expectation}. Got: {got}")] + Parse { + got: String, + expectation: &'static str, + }, +} + +#[cfg(all(test, unix))] +mod tests { + use super::*; + use pretty_assertions::assert_eq; + use std::process::{Command, Stdio}; + + #[test] + fn test_get_process_info() { + let subprocess = Command::new("bash") + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .expect("Can't execute command"); + let process_info = get_process_info(subprocess.id()); + let parent_pid = process_info.ok().and_then(|x| x.parent_pid); + assert_eq!(parent_pid, Some(std::process::id())); + } +} diff --git a/src/shell/infer/windows.rs b/src/shell/infer/windows.rs new file mode 100644 index 0000000..24dbf07 --- /dev/null +++ b/src/shell/infer/windows.rs @@ -0,0 +1,34 @@ +#![cfg(not(unix))] + +use crate::shell::Shell; +use log::debug; +use std::ffi::OsStr; +use sysinfo::{ProcessExt, System, SystemExt}; + +pub fn infer_shell() -> Option> { + let mut system = System::new(); + let mut current_pid = sysinfo::get_current_pid().ok(); + + while let Some(pid) = current_pid { + system.refresh_process(pid); + if let Some(process) = system.process(pid) { + current_pid = process.parent(); + let process_name = process + .exe() + .file_stem() + .and_then(OsStr::to_str) + .map(str::to_lowercase); + if let Some(shell) = process_name + .as_ref() + .map(|x| &x[..]) + .and_then(super::shell_from_string) + { + return Some(shell); + } + } else { + current_pid = None; + } + } + + None +}