use std::env;
use tokio::process::{Child, Command};
use crate::tools::{ToolError, ToolResult};
pub(crate) const TOOL_ENV_ALLOWLIST: &[&str] = &[
"PATH",
"PATHEXT",
"SystemRoot",
"SystemDrive",
"ComSpec",
"windir",
"USERPROFILE",
"HOMEDRIVE",
"HOMEPATH",
"APPDATA",
"LOCALAPPDATA",
"TEMP",
"TMP",
"NUMBER_OF_PROCESSORS",
"PROCESSOR_ARCHITECTURE",
"OS",
"PSModulePath",
"HOME",
"LANG",
"LC_ALL",
"LC_CTYPE",
"LC_MESSAGES",
"TZ",
"TERM",
"SHELL",
"TMPDIR",
"USER",
"LOGNAME",
"CARGO_HOME",
"RUSTUP_HOME",
"RUSTUP_TOOLCHAIN",
"PYTHONPATH",
"PYTHONHOME",
"NODE_PATH",
"GIT_EXEC_PATH",
"GIT_TEMPLATE_DIR",
];
fn is_allowed_env_name(name: &str) -> bool {
#[cfg(windows)]
{
TOOL_ENV_ALLOWLIST
.iter()
.any(|allowed| allowed.eq_ignore_ascii_case(name))
}
#[cfg(not(windows))]
{
TOOL_ENV_ALLOWLIST.contains(&name)
}
}
pub(crate) fn tool_child_env() -> Vec<(String, String)> {
env::vars()
.filter(|(name, _)| is_allowed_env_name(name))
.collect()
}
pub(crate) fn scrub_child_env(cmd: &mut Command) {
cmd.env_clear();
cmd.envs(tool_child_env());
}
pub(crate) struct GroupChild {
pub(crate) child: Child,
pub(crate) killer: GroupKiller,
}
impl GroupChild {
pub(crate) fn into_parts(self) -> (Child, GroupKiller) {
(self.child, self.killer)
}
}
pub(crate) struct GroupKiller {
#[cfg(windows)]
job: Option<win32job::Job>,
#[cfg(unix)]
pgid: Option<i32>,
disarmed: bool,
}
impl GroupKiller {
pub(crate) fn kill(&mut self) {
self.kill_impl();
self.disarmed = true;
}
pub(crate) fn disarm(&mut self) {
self.disarmed = true;
}
fn kill_impl(&mut self) {
#[cfg(windows)]
{
let _ = self.job.take();
}
#[cfg(unix)]
{
if let Some(pgid) = self.pgid.take() {
if let Some(pid) = rustix::process::Pid::from_raw(pgid) {
let _ = rustix::process::kill_process_group(pid, rustix::process::Signal::Kill);
}
}
}
}
}
impl Drop for GroupKiller {
fn drop(&mut self) {
if self.disarmed {
return;
}
self.kill_impl();
}
}
pub(crate) fn spawn_in_group(cmd: &mut Command) -> ToolResult<GroupChild> {
#[cfg(unix)]
{
cmd.process_group(0);
}
let child = cmd
.spawn()
.map_err(|e| ToolError::ExecutionError(format!("Failed to spawn process: {e}")))?;
#[cfg(windows)]
let killer = {
let assigned = match (build_kill_on_close_job(), child.raw_handle()) {
(Some(job), Some(handle)) => match job.assign_process(handle as isize) {
Ok(()) => Some(job),
Err(e) => {
warn_tree_kill_degraded(&format!("Job Object assignment failed: {e}"));
None
}
},
(None, _) => {
warn_tree_kill_degraded("Job Object could not be created");
None
}
(Some(_), None) => {
warn_tree_kill_degraded("child process handle was unavailable");
None
}
};
GroupKiller {
job: assigned,
disarmed: false,
}
};
#[cfg(unix)]
let killer = GroupKiller {
pgid: child.id().map(|id| id as i32),
disarmed: false,
};
#[cfg(not(any(windows, unix)))]
let killer = GroupKiller { disarmed: false };
Ok(GroupChild { child, killer })
}
#[cfg(windows)]
fn build_kill_on_close_job() -> Option<win32job::Job> {
let mut info = win32job::ExtendedLimitInfo::new();
info.limit_kill_on_job_close();
win32job::Job::create_with_limit_info(&info).ok()
}
#[cfg(windows)]
fn warn_tree_kill_degraded(reason: &str) {
eprintln!(
"warning: could not assign subprocess to a Windows Job Object; tree-kill \
degraded to direct-child-only (grandchildren may survive a --timeout kill): {reason}"
);
}
#[cfg(test)]
pub(crate) mod test_support {
pub(crate) const WORKER_SLEEP_SECS: u64 = 15;
pub(crate) const GRANDCHILD_PARENT_SLEEP_SECS: u64 = 18;
pub(crate) const CANCEL_FIRE_DELAY_MS: u64 = 5_000;
pub(crate) const POST_KILL_WAIT_MS: u64 = 13_000;
pub(crate) fn tree_kill_worker() -> String {
format!(
"\
import os, time\n\
d = os.path.dirname(os.path.abspath(__file__))\n\
open(os.path.join(d, 'start.marker'), 'w').close()\n\
time.sleep({WORKER_SLEEP_SECS})\n\
open(os.path.join(d, 'done.marker'), 'w').close()\n"
)
}
pub(crate) fn tree_kill_grandchild_worker() -> String {
format!(
"\
import os, sys, subprocess, time\n\
d = os.path.dirname(os.path.abspath(__file__))\n\
gc = (\n\
\"import os, time\\n\"\n\
\"d = \" + repr(d) + \"\\n\"\n\
\"open(os.path.join(d, 'start.marker'), 'w').close()\\n\"\n\
\"time.sleep({WORKER_SLEEP_SECS})\\n\"\n\
\"open(os.path.join(d, 'done.marker'), 'w').close()\\n\"\n\
)\n\
subprocess.Popen([sys.executable, \"-c\", gc])\n\
time.sleep({GRANDCHILD_PARENT_SLEEP_SECS})\n"
)
}
pub(crate) fn tree_kill_early_grandchild_worker() -> String {
format!(
"\
import os, sys, subprocess, time\n\
d = os.path.dirname(os.path.abspath(__file__))\n\
subprocess.Popen([sys.executable, \"-c\", \"import os, time\\n\" + \"d = \" + repr(d) + \"\\n\" + \"open(os.path.join(d, 'start.marker'), 'w').close()\\n\" + \"time.sleep({WORKER_SLEEP_SECS})\\n\" + \"open(os.path.join(d, 'done.marker'), 'w').close()\\n\"])\n\
time.sleep({GRANDCHILD_PARENT_SLEEP_SECS})\n"
)
}
pub(crate) const EXFIL_PROBE_WORKER: &str = "\
import os\n\
d = os.path.dirname(os.path.abspath(__file__))\n\
names = ['MAGI_PASSPHRASE', 'ANTHROPIC_API_KEY', 'OPENAI_API_KEY']\n\
found = [n for n in names if n in os.environ]\n\
raw = b''\n\
try: raw = open('/proc/self/environ', 'rb').read()\n\
except Exception: raw = b''\n\
found += [n + ':proc' for n in names if (n + '=').encode() in raw]\n\
open(os.path.join(d, 'exfil.marker'), 'w').write(','.join(found))\n";
pub(crate) fn python_available() -> bool {
std::process::Command::new("python")
.arg("--version")
.output()
.map(|o| o.status.success())
.unwrap_or(false)
}
}