Codex Handbook
shell-command/src/powershell.rs 272 lines
use std::path::PathBuf;use codex_utils_absolute_path::AbsolutePathBuf;use crate::command_safety::try_parse_powershell_ast_commands;use crate::shell_detect::ShellType;use crate::shell_detect::detect_shell_type;const POWERSHELL_FLAGS: &[&str] = &["-nologo", "-noprofile", "-command", "-c"];/// Prefixed command for powershell shell calls to request UTF-8 console output.pub const UTF8_OUTPUT_PREFIX: &str =    "try { [Console]::OutputEncoding=[System.Text.Encoding]::UTF8 } catch {}\n";pub fn prefix_powershell_script_with_utf8(command: &[String]) -> Vec<String> {    let Some((_, script)) = extract_powershell_command(command) else {        return command.to_vec();    };    let trimmed = script.trim_start();    let script = if trimmed.starts_with(UTF8_OUTPUT_PREFIX) {        script.to_string()    } else {        format!("{UTF8_OUTPUT_PREFIX}{script}")    };    let mut command: Vec<String> = command[..(command.len() - 1)]        .iter()        .map(std::string::ToString::to_string)        .collect();    command.push(script);    command}/// Extract the PowerShell script body from an invocation such as:////// - ["pwsh", "-NoProfile", "-Command", "Get-ChildItem -Recurse | Select-String foo"]/// - ["powershell.exe", "-Command", "Write-Host hi"]/// - ["powershell", "-NoLogo", "-NoProfile", "-Command", "...script..."]////// Returns (`shell`, `script`) when the first arg is a PowerShell executable and a/// `-Command` (or `-c`) flag is present followed by a script string.pub fn extract_powershell_command(command: &[String]) -> Option<(&str, &str)> {    if command.len() < 3 {        return None;    }    let shell = &command[0];    if !matches!(        detect_shell_type(PathBuf::from(shell)),        Some(ShellType::PowerShell)    ) {        return None;    }    // Find the first occurrence of -Command (accept common short alias -c as well)    let mut i = 1usize;    while i + 1 < command.len() {        let flag = &command[i];        // Reject unknown flags        if !POWERSHELL_FLAGS.contains(&flag.to_ascii_lowercase().as_str()) {            return None;        }        if flag.eq_ignore_ascii_case("-Command") || flag.eq_ignore_ascii_case("-c") {            let script = &command[i + 1];            return Some((shell, script));        }        i += 1;    }    None}/// Parse the script body from a top-level PowerShell wrapper into argv-like commands.////// This is intentionally narrower than the Windows safe-command parser: it only unwraps the/// `-Command`/`-c` body from a PowerShell invocation we already recognize, then delegates the/// script itself to the PowerShell AST parser.pub fn parse_powershell_command_into_plain_commands(    command: &[String],) -> Option<Vec<Vec<String>>> {    let (executable, script) = extract_powershell_command(command)?;    try_parse_powershell_ast_commands(executable, script)}/// This function attempts to find a powershell.exe executable on the system.pub fn try_find_powershell_executable_blocking() -> Option<AbsolutePathBuf> {    try_find_powershellish_executable_in_path(&["powershell.exe"])}/// This function attempts to find a pwsh.exe executable on the system./// Note that pwsh.exe and powershell.exe are different executables:////// - pwsh.exe is the cross-platform PowerShell Core (v6+) executable/// - powershell.exe is the Windows PowerShell (v5.1 and earlier) executable////// Further, while powershell.exe is included by default on Windows systems,/// pwsh.exe must be installed separately by the user. And even when the user/// has installed pwsh.exe, it may not be available in the system PATH, in which/// case we attempt to locate it via other means.pub fn try_find_pwsh_executable_blocking() -> Option<AbsolutePathBuf> {    if let Some(ps_home) = std::process::Command::new("cmd")        .args(["/C", "pwsh", "-NoProfile", "-Command", "$PSHOME"])        .output()        .ok()        .and_then(|out| {            if !out.status.success() {                return None;            }            let stdout = String::from_utf8_lossy(&out.stdout);            let trimmed = stdout.trim();            (!trimmed.is_empty()).then(|| trimmed.to_string())        })    {        let candidate = AbsolutePathBuf::resolve_path_against_base("pwsh.exe", &ps_home);        if is_powershellish_executable_available(candidate.as_path()) {            return Some(candidate);        }    }    try_find_powershellish_executable_in_path(&["pwsh.exe"])}fn try_find_powershellish_executable_in_path(candidates: &[&str]) -> Option<AbsolutePathBuf> {    for candidate in candidates {        let Ok(resolved_path) = which::which(candidate) else {            continue;        };        if !is_powershellish_executable_available(&resolved_path) {            continue;        }        let Ok(abs_path) = AbsolutePathBuf::from_absolute_path(resolved_path) else {            continue;        };        return Some(abs_path);    }    None}fn is_powershellish_executable_available(powershell_or_pwsh_exe: &std::path::Path) -> bool {    // This test works for both powershell.exe and pwsh.exe.    std::process::Command::new(powershell_or_pwsh_exe)        .args(["-NoLogo", "-NoProfile", "-Command", "Write-Output ok"])        .output()        .map(|output| output.status.success())        .unwrap_or(false)}#[cfg(test)]mod tests {    use super::UTF8_OUTPUT_PREFIX;    use super::extract_powershell_command;    #[cfg(windows)]    use super::parse_powershell_command_into_plain_commands;    use super::prefix_powershell_script_with_utf8;    #[test]    fn extracts_basic_powershell_command() {        let cmd = vec![            "powershell".to_string(),            "-Command".to_string(),            "Write-Host hi".to_string(),        ];        let (_shell, script) = extract_powershell_command(&cmd).expect("extract");        assert_eq!(script, "Write-Host hi");    }    #[test]    fn extracts_lowercase_flags() {        let cmd = vec![            "powershell".to_string(),            "-nologo".to_string(),            "-command".to_string(),            "Write-Host hi".to_string(),        ];        let (_shell, script) = extract_powershell_command(&cmd).expect("extract");        assert_eq!(script, "Write-Host hi");    }    #[test]    fn extracts_full_path_powershell_command() {        let command = if cfg!(windows) {            "C:\\windows\\System32\\WindowsPowerShell\\v1.0\\powershell.exe".to_string()        } else {            "/usr/local/bin/powershell.exe".to_string()        };        let cmd = vec![command, "-Command".to_string(), "Write-Host hi".to_string()];        let (_shell, script) = extract_powershell_command(&cmd).expect("extract");        assert_eq!(script, "Write-Host hi");    }    #[test]    fn extracts_with_noprofile_and_alias() {        let cmd = vec![            "pwsh".to_string(),            "-NoProfile".to_string(),            "-c".to_string(),            "Get-ChildItem | Select-String foo".to_string(),        ];        let (_shell, script) = extract_powershell_command(&cmd).expect("extract");        assert_eq!(script, "Get-ChildItem | Select-String foo");    }    #[test]    fn prefixes_powershell_command_with_best_effort_utf8() {        let cmd = vec![            "powershell".to_string(),            "-Command".to_string(),            "Write-Host hi".to_string(),        ];        let prefixed = prefix_powershell_script_with_utf8(&cmd);        assert_eq!(            prefixed,            vec![                "powershell".to_string(),                "-Command".to_string(),                format!("{UTF8_OUTPUT_PREFIX}Write-Host hi"),            ]        );    }    #[test]    fn does_not_duplicate_utf8_prefix() {        let cmd = vec![            "powershell".to_string(),            "-Command".to_string(),            format!("{UTF8_OUTPUT_PREFIX}Write-Host hi"),        ];        assert_eq!(prefix_powershell_script_with_utf8(&cmd), cmd);    }    #[cfg(windows)]    #[test]    fn parses_plain_powershell_commands() {        let commands = parse_powershell_command_into_plain_commands(&[            "powershell.exe".to_string(),            "-NoProfile".to_string(),            "-Command".to_string(),            "echo hi".to_string(),        ])        .expect("parse");        assert_eq!(commands, vec![vec!["echo".to_string(), "hi".to_string()]]);    }    #[cfg(windows)]    #[test]    fn parses_multiple_plain_powershell_commands() {        let commands = parse_powershell_command_into_plain_commands(&[            "powershell.exe".to_string(),            "-NoProfile".to_string(),            "-Command".to_string(),            "Write-Output foo | Measure-Object".to_string(),        ])        .expect("parse");        assert_eq!(            commands,            vec![                vec!["Write-Output".to_string(), "foo".to_string()],                vec!["Measure-Object".to_string()],            ]        );    }}