diff --git a/src/terminal.rs b/src/terminal.rs index 48a3b35..0b7ffcc 100644 --- a/src/terminal.rs +++ b/src/terminal.rs @@ -1,6 +1,12 @@ use std::env; use sysinfo::{Pid, PidExt, ProcessExt, System, SystemExt}; +// Allows detection of shells that host themselves (i.e. Command Prompt). +const SELF_HOSTED_SHELLS: [&str; 2] = [ + "cmd.exe", + "powershell.exe" +]; + macro_rules! env_to_name { ($env: expr, $name: expr) => { match env::var($env) { @@ -24,6 +30,7 @@ pub fn get_terminal() -> Option { Err(_) => (), }; + env_to_name!("ConEmuPID", "ConEmu"); env_to_name!("WT_SESSION", "Windows Terminal"); env_to_name!("SSH_CONNECTION", "SSH"); env_to_name!("tw52", "TosWin2"); @@ -32,12 +39,16 @@ pub fn get_terminal() -> Option { let mut pid = Pid::from_u32(std::process::id()); let shell = match env::var("SHELL") { Ok(f) => f, - Err(_) => return None, + Err(_) => "".to_string(), }; let shell_name = shell.split("/").last().unwrap(); + let mut name: Option = None; + + let mut self_hosted = false; + // Get to the shell PID. - loop { + 'find_shell: loop { let ppid = pid_to_ppid(pid); if ppid.is_none() { // We ran out of parents. @@ -48,32 +59,46 @@ pub fn get_terminal() -> Option { pid = ppid.unwrap(); // It should be safe to unwrap the PPID now. - let name = pid_to_name(pid); + let pid_name = pid_to_name(pid); - if name.is_none() { + if pid_name.is_none() { continue; } - println!("{}", name.as_ref().unwrap()); + let name_unwrapped = pid_name.unwrap(); - let name_unwrapped = name.unwrap(); + // Detect self-hosted shells. + for shell in SELF_HOSTED_SHELLS { + if name_unwrapped == shell { + self_hosted = true; + name = Some(name_unwrapped); + break 'find_shell; + } + } if name_unwrapped == shell.as_str() || name_unwrapped == shell_name { + name = Some(name_unwrapped); break; // We found the shell. } } - // Try to get parent once more. - pid = match pid_to_ppid(pid) { - None => return None, - Some(f) => f, - }; + if !self_hosted { + // Try to get parent once more. + match pid_to_ppid(pid) { + None => return None, + Some(f) => { + // Try to get name. + name = pid_to_name(f); + }, + }; + } - // Get the name. - return match pid_to_name(pid) { + return match name { + Some(f) => { + Some(f.split(".").nth(0).unwrap().to_string()) + }, None => None, - Some(f) => Some(f), - }; + } } fn pid_to_name(pid: Pid) -> Option {