Skip to content

Commit eac8a8a

Browse files
author
Deimos Agent
committed
fix: safe-env whitelist in LocalInteractiveSession; extra_env parity in SSHInteractiveSession
Addresses two concerns raised by @frdel in PR review: 1. shell_local.py — replace os.environ merge with _SAFE_ENV_KEYS whitelist Previously: env = {**os.environ, **self.extra_env} if self.extra_env else None This leaked all framework env vars (API keys, tokens) into the subprocess. Now: only PATH, HOME, USER, SHELL, TERM, LANG, LC_ALL, TMPDIR, PWD are forwarded from the host environment; extra_env values are merged on top. 2. shell_ssh.py — add extra_env: dict | None = None parameter to SSHInteractiveSession.__init__ matching shell_local.py signature. Extra vars are injected via 'export KEY=VALUE' in the initial_command block (shlex.quote-escaped) so they are available session-wide. Paramiko invoke_shell() does not pass env to the server reliably (AcceptEnv restrictions), so the export-prefix approach is used. Both classes now have identical extra_env: dict | None = None signatures. No call-site changes required — extra_env defaults to None (no behaviour change for existing users).
1 parent d357c24 commit eac8a8a

2 files changed

Lines changed: 40 additions & 11 deletions

File tree

plugins/_code_execution/helpers/shell_local.py

Lines changed: 19 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import os
12
import platform
23
import select
34
import subprocess
@@ -8,14 +9,27 @@
89
from plugins._code_execution.helpers import tty_session
910
from plugins._code_execution.helpers.shell_ssh import clean_string
1011

12+
# Environment variable keys that are safe to forward to the subprocess.
13+
# Deliberately excludes API keys, bearer tokens, secrets, and all framework
14+
# env vars — only the minimal set required for a functional interactive shell.
15+
_SAFE_ENV_KEYS = {"PATH", "HOME", "USER", "SHELL", "TERM", "LANG", "LC_ALL", "TMPDIR", "PWD"}
16+
17+
1118
class LocalInteractiveSession:
12-
def __init__(self, cwd: str|None = None):
13-
self.session: tty_session.TTYSession|None = None
19+
def __init__(self, cwd: str | None = None, extra_env: dict | None = None):
20+
self.session: tty_session.TTYSession | None = None
1421
self.full_output = ''
1522
self.cwd = cwd
23+
self.extra_env = extra_env
1624

1725
async def connect(self):
18-
self.session = tty_session.TTYSession(runtime.get_terminal_executable(), cwd=self.cwd)
26+
# When extra_env is provided, build a clean env from the safe whitelist
27+
# only — never merge the full os.environ, which would expose API keys
28+
# and other framework secrets to the subprocess.
29+
env = (
30+
{k: v for k, v in os.environ.items() if k in _SAFE_ENV_KEYS} | self.extra_env
31+
) if self.extra_env else None
32+
self.session = tty_session.TTYSession(runtime.get_terminal_executable(), cwd=self.cwd, env=env)
1933
await self.session.start()
2034
await self.session.read_full_until_idle(idle_timeout=1, total_timeout=1)
2135

@@ -29,7 +43,7 @@ async def send_command(self, command: str):
2943
raise Exception("Shell not connected")
3044
self.full_output = ""
3145
await self.session.sendline(command)
32-
46+
3347
async def read_output(self, timeout: float = 0, reset_full_output: bool = False) -> Tuple[str, Optional[str]]:
3448
if not self.session:
3549
raise Exception("Shell not connected")
@@ -47,4 +61,4 @@ async def read_output(self, timeout: float = 0, reset_full_output: bool = False)
4761

4862
if not partial_output:
4963
return clean_full_output, None
50-
return clean_full_output, partial_output
64+
return clean_full_output, partial_output

plugins/_code_execution/helpers/shell_ssh.py

Lines changed: 21 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import asyncio
22
import paramiko
3+
import shlex
34
import time
45
import re
56
from typing import Tuple
@@ -14,7 +15,8 @@ class SSHInteractiveSession:
1415
# ps1_label = "SSHInteractiveSession CLI>"
1516

1617
def __init__(
17-
self, logger: Log, hostname: str, port: int, username: str, password: str, cwd: str|None = None
18+
self, logger: Log, hostname: str, port: int, username: str, password: str,
19+
cwd: str | None = None, extra_env: dict | None = None
1820
):
1921
self.logger = logger
2022
self.hostname = hostname
@@ -28,6 +30,7 @@ def __init__(
2830
self.last_command = b""
2931
self.trimmed_command_length = 0 # Initialize trimmed_command_length
3032
self.cwd = cwd
33+
self.extra_env = extra_env
3134

3235
async def connect(self, keepalive_interval: int = 5):
3336
"""
@@ -37,7 +40,7 @@ async def connect(self, keepalive_interval: int = 5):
3740
----------
3841
keepalive_interval : int
3942
Interval in **seconds** between keep-alive packets sent by Paramiko.
40-
A value ≤ 0 disables Paramikos keep-alive feature.
43+
A value ≤ 0 disables Paramiko's keep-alive feature.
4144
"""
4245
errors = 0
4346
while True:
@@ -66,6 +69,17 @@ async def connect(self, keepalive_interval: int = 5):
6669
initial_command = "unset PROMPT_COMMAND PS0; stty -echo"
6770
if self.cwd:
6871
initial_command = f"cd {self.cwd}; {initial_command}"
72+
73+
# When extra_env is provided, prepend export statements so the
74+
# variables are available for the entire session. Values are
75+
# shell-quoted via shlex.quote to prevent injection.
76+
if self.extra_env:
77+
exports = "; ".join(
78+
f"export {k}={shlex.quote(str(v))}"
79+
for k, v in self.extra_env.items()
80+
)
81+
initial_command = f"{exports}; {initial_command}"
82+
6983
self.shell.send(f"{initial_command}\n".encode())
7084

7185
# wait for initial prompt/output to settle
@@ -104,7 +118,7 @@ async def send_command(self, command: str):
104118
self.last_command = command.encode()
105119
self.trimmed_command_length = 0
106120
self.shell.send(self.last_command)
107-
121+
108122
async def read_output(
109123
self, timeout: float = 0, reset_full_output: bool = False
110124
) -> Tuple[str, str]:
@@ -138,7 +152,7 @@ async def read_output(
138152
# deviation_threshold=8,
139153
# deviation_reset=2,
140154
# ignore_patterns=[
141-
# rb"\x1b\[\?\d{4}[a-zA-Z](?:> )?", # ANSI escape sequences
155+
# rb"\[\?\d{4}[a-zA-Z](?:> )?", # ANSI escape sequences
142156
# rb"\r", # Carriage return
143157
# rb">\s", # Greater-than symbol
144158
# ],
@@ -212,13 +226,14 @@ def recv_all(num_bytes):
212226

213227
return data
214228

229+
215230
def clean_string(input_string):
216231
# Remove ANSI escape codes
217-
ansi_escape = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
232+
ansi_escape = re.compile(r"(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
218233
cleaned = ansi_escape.sub("", input_string)
219234

220235
# remove null bytes
221-
cleaned = cleaned.replace("\x00", "")
236+
cleaned = cleaned.replace("", "")
222237

223238
# remove ipython \r\r\n> sequences from the start
224239
cleaned = re.sub(r'^[ \r]*(?:\r*\n>[ \r]*)*', '', cleaned)

0 commit comments

Comments
 (0)