magi-code 0.77.1

use crate::security::is_credential_like_key;
use std::{
    collections::BTreeMap,
    ffi::{OsStr, OsString},
    process::Command,
};

/// The explicit environment contract for a deny-by-default child process.
///
/// These profiles are intentionally separate: a child receives only the ambient variables
/// declared by its profile, plus the portable Windows startup baseline where that profile needs
/// it. Values supplied explicitly by MCP configuration are applied after the ambient filter.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum SubprocessEnvProfile {
    /// Hook and other sanitized shell execution: shell locale and identity variables.
    Shell,
    /// LSP servers intentionally receive PATH only.
    Lsp,
    /// MCP stdio servers receive PATH and the portable baseline before configured values.
    McpStdio,
}

#[derive(Debug, Clone, Copy)]
struct ProfileRules {
    exact: &'static [&'static str],
    locale: &'static [&'static str],
    include_windows_startup: bool,
}

const EMPTY: &[&str] = &[];
const LOCALE_VARIABLES: &[&str] = &[
    "LC_CTYPE",
    "LC_NUMERIC",
    "LC_TIME",
    "LC_COLLATE",
    "LC_MONETARY",
    "LC_MESSAGES",
    "LC_PAPER",
    "LC_NAME",
    "LC_ADDRESS",
    "LC_TELEPHONE",
    "LC_MEASUREMENT",
    "LC_IDENTIFICATION",
    "LC_ALL",
];
const SHELL_VARIABLES: &[&str] = &["PATH", "HOME", "USER", "LOGNAME", "SHELL", "TMPDIR", "LANG"];
const PATH_ONLY: &[&str] = &["PATH"];

#[cfg(windows)]
fn matches_environment_name(actual: &str, expected: &str) -> bool {
    actual.eq_ignore_ascii_case(expected)
}

#[cfg(not(windows))]
fn matches_environment_name(actual: &str, expected: &str) -> bool {
    actual == expected
}

impl SubprocessEnvProfile {
    fn rules(self) -> ProfileRules {
        match self {
            Self::Shell => ProfileRules {
                exact: SHELL_VARIABLES,
                locale: LOCALE_VARIABLES,
                include_windows_startup: true,
            },
            Self::Lsp => ProfileRules {
                exact: PATH_ONLY,
                locale: EMPTY,
                include_windows_startup: false,
            },
            Self::McpStdio => ProfileRules {
                exact: PATH_ONLY,
                locale: EMPTY,
                include_windows_startup: true,
            },
        }
    }

    /// Returns whether an ambient variable belongs to this profile.
    ///
    /// Credential-shaped names are rejected even when a future profile allowlist could otherwise
    /// match them. Explicit MCP configuration is deliberately handled separately and is not
    /// treated as ambient input.
    pub(crate) fn allows_ambient(self, key: &OsStr) -> bool {
        let Some(key) = key.to_str() else {
            return false;
        };
        if is_credential_like_key(key) {
            return false;
        }

        let rules = self.rules();
        (rules.include_windows_startup && is_windows_startup_variable(key))
            || rules
                .exact
                .iter()
                .any(|expected| matches_environment_name(key, expected))
            || rules
                .locale
                .iter()
                .any(|expected| matches_environment_name(key, expected))
    }
}

/// Clear inherited variables and apply one explicit ambient environment profile.
pub(crate) fn apply_profile(command: &mut Command, profile: SubprocessEnvProfile) {
    apply_profile_from(command, profile, std::env::vars_os());
}

/// Apply a profile, then overlay explicitly configured MCP values.
///
/// Configured values intentionally have precedence over ambient PATH and platform startup values;
/// this preserves the MCP `env` contract without reopening ambient inheritance.
pub(crate) fn apply_configured_profile(
    command: &mut Command,
    profile: SubprocessEnvProfile,
    configured: &BTreeMap<String, String>,
) {
    apply_configured_profile_from(command, profile, std::env::vars_os(), configured);
}

fn apply_profile_from(
    command: &mut Command,
    profile: SubprocessEnvProfile,
    environment: impl IntoIterator<Item = (OsString, OsString)>,
) {
    command.env_clear();
    for (key, value) in environment {
        if profile.allows_ambient(&key) {
            command.env(key, value);
        }
    }
}

fn apply_configured_profile_from(
    command: &mut Command,
    profile: SubprocessEnvProfile,
    environment: impl IntoIterator<Item = (OsString, OsString)>,
    configured: &BTreeMap<String, String>,
) {
    apply_profile_from(command, profile, environment);
    command.envs(configured);
}

#[cfg(windows)]
fn is_windows_startup_variable(key: &str) -> bool {
    ["ComSpec", "PATHEXT", "SystemRoot", "WINDIR"]
        .iter()
        .any(|name| matches_environment_name(key, name))
}

#[cfg(not(windows))]
fn is_windows_startup_variable(_key: &str) -> bool {
    false
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::BTreeMap;

    fn environment(entries: &[(&str, &str)]) -> Vec<(OsString, OsString)> {
        entries
            .iter()
            .map(|(key, value)| (OsString::from(key), OsString::from(value)))
            .collect()
    }

    fn command_environment(command: &Command) -> BTreeMap<String, String> {
        command
            .get_envs()
            .filter_map(|(key, value)| {
                value.map(|value| {
                    (
                        key.to_string_lossy().into_owned(),
                        value.to_string_lossy().into_owned(),
                    )
                })
            })
            .collect()
    }

    fn all_profiles() -> [SubprocessEnvProfile; 3] {
        [
            SubprocessEnvProfile::Shell,
            SubprocessEnvProfile::Lsp,
            SubprocessEnvProfile::McpStdio,
        ]
    }

    #[test]
    fn profiles_keep_distinct_allowlists() {
        assert!(SubprocessEnvProfile::Shell.allows_ambient(OsStr::new("USER")));
        assert!(SubprocessEnvProfile::Shell.allows_ambient(OsStr::new("SHELL")));
        assert!(!SubprocessEnvProfile::Shell.allows_ambient(OsStr::new("DISPLAY")));
        assert!(SubprocessEnvProfile::Shell.allows_ambient(OsStr::new("LC_ALL")));
        assert!(!SubprocessEnvProfile::Lsp.allows_ambient(OsStr::new("HOME")));
        assert!(SubprocessEnvProfile::McpStdio.allows_ambient(OsStr::new("PATH")));
    }

    #[test]
    fn exact_and_windows_startup_names_follow_platform_case_rules() {
        assert!(SubprocessEnvProfile::Shell.allows_ambient(OsStr::new("PATH")));
        assert_eq!(
            SubprocessEnvProfile::Shell.allows_ambient(OsStr::new("Path")),
            cfg!(windows)
        );
        assert_eq!(
            SubprocessEnvProfile::Shell.allows_ambient(OsStr::new("path")),
            cfg!(windows)
        );

        for name in ["SystemRoot", "systemroot", "SYSTEMROOT"] {
            assert_eq!(
                SubprocessEnvProfile::Shell.allows_ambient(OsStr::new(name)),
                cfg!(windows),
                "unexpected startup-variable match for {name}"
            );
        }
    }

    #[test]
    fn locale_allowlist_uses_exact_standard_names_and_platform_case_rules() {
        let locale_profiles = [SubprocessEnvProfile::Shell];
        for profile in locale_profiles {
            for name in [
                "LANG",
                "LC_CTYPE",
                "LC_NUMERIC",
                "LC_TIME",
                "LC_COLLATE",
                "LC_MONETARY",
                "LC_MESSAGES",
                "LC_PAPER",
                "LC_NAME",
                "LC_ADDRESS",
                "LC_TELEPHONE",
                "LC_MEASUREMENT",
                "LC_IDENTIFICATION",
                "LC_ALL",
            ] {
                assert!(
                    profile.allows_ambient(OsStr::new(name)),
                    "{profile:?} rejected standard locale variable {name}"
                );
            }
        }

        for name in ["LC_PASSWORD", "LC_PRIVATE_KEY", "LC_API_TOKEN"] {
            for profile in all_profiles() {
                assert!(
                    !profile.allows_ambient(OsStr::new(name)),
                    "{profile:?} accepted non-standard locale variable {name}"
                );
            }
        }

        for name in ["lc_ALL", "Lc_All", "lC_ctype"] {
            assert_eq!(
                SubprocessEnvProfile::Shell.allows_ambient(OsStr::new(name)),
                cfg!(windows),
                "unexpected locale-name case match for {name}"
            );
        }
        assert!(!SubprocessEnvProfile::Shell.allows_ambient(OsStr::new("LC")));
    }

    #[test]
    fn credential_like_ambient_names_are_rejected_for_every_profile() {
        for profile in all_profiles() {
            for name in [
                "OPENAI_API_KEY",
                "MC_API_KEY",
                "ANTHROPIC_API_KEY",
                "AWS_SECRET_ACCESS_KEY",
                "MCP_TOKEN",
                "LC_API_KEY",
            ] {
                assert!(
                    !profile.allows_ambient(OsStr::new(name)),
                    "{profile:?} accepted ambient credential {name}"
                );
            }
        }
    }

    #[test]
    fn windows_startup_baseline_is_profile_specific_and_portable() {
        let baseline = ["ComSpec", "PATHEXT", "SystemRoot", "WINDIR"];
        for name in baseline {
            #[cfg(windows)]
            {
                assert!(SubprocessEnvProfile::Shell.allows_ambient(OsStr::new(name)));
                assert!(SubprocessEnvProfile::McpStdio.allows_ambient(OsStr::new(name)));
                assert!(!SubprocessEnvProfile::Lsp.allows_ambient(OsStr::new(name)));
            }
            #[cfg(not(windows))]
            {
                assert!(!SubprocessEnvProfile::Shell.allows_ambient(OsStr::new(name)));
            }
        }
    }

    #[test]
    fn apply_profile_clears_existing_values_and_applies_only_selected_ambient_values() {
        let mut command = Command::new("profile-test");
        command.env("STALE_VALUE", "must disappear");
        apply_profile_from(
            &mut command,
            SubprocessEnvProfile::Shell,
            environment(&[
                ("PATH", "path"),
                ("HOME", "home"),
                ("TMPDIR", "tmp"),
                ("LANG", "lang"),
                ("LC_ALL", "locale"),
                ("LC_API_KEY", "credential"),
                ("OPENAI_API_KEY", "credential"),
                ("DISPLAY", "display"),
            ]),
        );

        let actual = command_environment(&command);
        assert_eq!(actual.get("PATH"), Some(&"path".to_string()));
        assert_eq!(actual.get("HOME"), Some(&"home".to_string()));
        assert_eq!(actual.get("LC_ALL"), Some(&"locale".to_string()));
        assert!(!actual.contains_key("STALE_VALUE"));
        assert!(!actual.contains_key("LC_API_KEY"));
        assert!(!actual.contains_key("OPENAI_API_KEY"));
        assert!(!actual.contains_key("DISPLAY"));
    }

    #[test]
    fn configured_mcp_values_override_ambient_values_after_filtering() {
        let mut command = Command::new("mcp-profile-test");
        let configured = BTreeMap::from([
            ("PATH".to_string(), "configured-path".to_string()),
            ("MCP_TOKEN".to_string(), "configured-token".to_string()),
            ("MCP_ROOT".to_string(), "/configured/root".to_string()),
        ]);
        apply_configured_profile_from(
            &mut command,
            SubprocessEnvProfile::McpStdio,
            environment(&[
                ("PATH", "ambient-path"),
                ("MCP_TOKEN", "ambient-token"),
                ("OPENAI_API_KEY", "ambient-key"),
                ("HOME", "ambient-home"),
            ]),
            &configured,
        );

        let actual = command_environment(&command);
        assert_eq!(actual.get("PATH"), Some(&"configured-path".to_string()));
        assert_eq!(
            actual.get("MCP_TOKEN"),
            Some(&"configured-token".to_string())
        );
        assert_eq!(
            actual.get("MCP_ROOT"),
            Some(&"/configured/root".to_string())
        );
        assert!(!actual.contains_key("OPENAI_API_KEY"));
        assert!(!actual.contains_key("HOME"));
    }
}