"""Tests for the restricted shell mode of the MCP bridge: python3 -m unittest scripts/test_mcp_shell.py """ import json import os import pathlib import shutil import sys import tempfile import threading import time import unittest HERE = pathlib.Path(__file__).resolve().parent sys.path.insert(0, str(HERE)) import mcp_bridge # noqa: E402 import mcp_shell # noqa: E402 class FakeConfirm: """Answers every prompt with `answer` and remembers what it was shown.""" def __init__(self, answer=True): self.answer, self.prompts = answer, [] def __call__(self, text): self.prompts.append(text) return (self.answer, None if self.answer else '终端里拒绝了这条命令') def call(session, command, stdin=None, mid=1): args = {'command': command, **({'stdin': stdin} if stdin is not None else {})} reply = json.loads(session.rpc({'jsonrpc': '2.0', 'id': mid, 'method': 'tools/call', 'params': {'name': 'run', 'arguments': args}}, '')) return reply['result'] def text(result): return '\n'.join(c['text'] for c in result['content']) class ShellSessionTest(unittest.TestCase): def setUp(self): self.tmp = tempfile.TemporaryDirectory() self.addCleanup(self.tmp.cleanup) self.base = pathlib.Path(self.tmp.name) (self.base / 'hello.txt').write_text('hello\nworld\n', encoding='utf-8') self.confirm = FakeConfirm(True) def session(self, **shell): cfg = {'id': 'sh', 'name': '受限 shell', 'kind': 'shell', 'shell': {'allow': ['cat', 'wc', 'echo', 'printf'], **shell}} return mcp_shell.ShellSession(cfg, self.base, mcp_bridge.Events(), confirm=self.confirm) def test_lists_one_run_tool_naming_the_allowed_programs(self): s = self.session() init = json.loads(s.rpc({'jsonrpc': '2.0', 'id': 1, 'method': 'initialize', 'params': {'protocolVersion': '2025-03-26'}}, '')) self.assertEqual(init['result']['protocolVersion'], '2025-03-26') tools = json.loads(s.rpc({'jsonrpc': '2.0', 'id': 2, 'method': 'tools/list'}, ''))['result']['tools'] self.assertEqual([t['name'] for t in tools], ['run']) self.assertIn('cat', tools[0]['description']) self.assertEqual(tools[0]['inputSchema']['required'], ['command']) def test_allowed_command_runs_after_confirmation(self): result = call(self.session(), 'wc -l hello.txt') self.assertFalse(result['isError'], text(result)) self.assertIn('2 hello.txt', text(result)) self.assertEqual(len(self.confirm.prompts), 1) self.assertIn('wc -l hello.txt', self.confirm.prompts[0]) def test_stdin_is_passed_to_the_program(self): self.assertIn('3', text(call(self.session(), 'wc -l', stdin='a\nb\nc\n'))) def test_refused_confirmation_does_not_run(self): self.confirm.answer = False marker = self.base / 'ran' result = call(self.session(allow=['touch']), f'touch {marker}') self.assertTrue(result['isError']) self.assertIn('拒绝', text(result)) self.assertFalse(marker.exists()) def test_programs_outside_the_allowlist_are_refused_before_asking(self): for command in ('rm hello.txt', 'sh -c "cat hello.txt"', f'{shutil.which("rm")} hello.txt', 'env cat hello.txt'): result = call(self.session(), command) self.assertTrue(result['isError'], command) self.assertIn('不在白名单', text(result), command) self.assertEqual(self.confirm.prompts, []) self.assertTrue((self.base / 'hello.txt').exists()) def test_a_same_named_program_elsewhere_is_not_the_allowed_one(self): fake = self.base / 'cat' fake.write_text('#!/bin/sh\necho pwned\n') fake.chmod(0o755) for command in ('./cat hello.txt', f'{fake} hello.txt'): self.assertIn('不在白名单', text(call(self.session(), command)), command) def test_the_full_path_of_an_allowed_program_is_fine(self): result = call(self.session(), f'{shutil.which("cat")} hello.txt') self.assertFalse(result['isError'], text(result)) def test_shell_syntax_is_just_text(self): # No shell: pipes, redirects, substitutions and globs reach the program as literal words result = call(self.session(), 'echo $(id) `id` ; rm -rf x | cat > out *.txt') self.assertEqual(text(result).splitlines()[0], '$(id) `id` ; rm -rf x | cat > out *.txt') self.assertFalse((self.base / 'out').exists()) def test_bad_quoting_and_empty_commands_are_errors(self): self.assertIn('引号', text(call(self.session(), 'echo "unclosed'))) self.assertIn('空', text(call(self.session(), ' '))) self.assertEqual(self.confirm.prompts, []) def test_confirm_false_skips_the_prompt(self): result = call(self.session(confirm=False), 'echo hi') self.assertEqual(text(result).splitlines()[0], 'hi') self.assertEqual(self.confirm.prompts, []) def test_prompt_escapes_control_characters(self): # An escape sequence could otherwise rewrite the terminal line the user is reading call(self.session(), "printf '\x1b[2K\x1b[1Gls\r'") prompt = self.confirm.prompts[0] self.assertNotIn('\x1b', prompt) self.assertNotIn('\r', prompt) self.assertIn('\\x1b', prompt) def test_newlines_and_tabs_cannot_fake_a_prompt(self): # A quoted argument full of newlines could otherwise push the real command off-screen # and print a harmless-looking fake question right above [y/N] call(self.session(), "echo 'safe\n\n\n\n[MCP 受限 shell] 命令：ls\t允许运行吗？'") prompt = self.confirm.prompts[0] command_line = [l for l in prompt.splitlines() if l.startswith(' 命令：')] self.assertEqual(len(command_line), 1) self.assertIn('\\x0a', command_line[0]) self.assertIn('\\x09', command_line[0]) self.assertEqual(len(prompt.splitlines()), 4, prompt) def test_over_long_commands_are_refused_before_asking(self): result = call(self.session(), 'echo ' + 'x' * mcp_shell.MAX_COMMAND) self.assertIn('太长', text(result)) self.assertEqual(self.confirm.prompts, []) def test_truncated_stdin_is_called_out(self): call(self.session(), 'wc -c', stdin='a' * 1000) self.assertIn('没有显示', self.confirm.prompts[0]) def test_program_replaced_while_waiting_is_not_run(self): tool = self.base / 'tool.sh' tool.write_text('#!/bin/sh\necho original\n') tool.chmod(0o755) def swap_during_confirmation(text): swapped = self.base / 'swapped.sh' swapped.write_text('#!/bin/sh\necho swapped\n') swapped.chmod(0o755) os.replace(swapped, tool) return True, None self.confirm = swap_during_confirmation result = call(self.session(allow=['./tool.sh']), './tool.sh') self.assertTrue(result['isError']) self.assertIn('被替换', text(result)) def test_prompt_escapes_direction_overrides(self): call(self.session(), 'echo "\u202eetadpu"') self.assertNotIn('\u202e', self.confirm.prompts[0]) self.assertIn('\\u202e', self.confirm.prompts[0]) def test_timeout_and_output_limits_apply(self): result = call(self.session(allow=['sleep'], timeout_s=0.3), 'sleep 5') self.assertIn('超过', text(result)) class ConfigTest(unittest.TestCase): def load(self, servers): with tempfile.NamedTemporaryFile('w', suffix='.json', delete=False) as f: json.dump({'servers': servers}, f) try: return mcp_bridge.load_config(f.name)[0] finally: pathlib.Path(f.name).unlink() def test_shell_servers_load(self): servers = self.load([{'id': 'sh', 'shell': {'allow': ['ls', 'git'], 'timeout_s': 5}}]) self.assertEqual(servers['sh']['kind'], 'shell') def test_server_list_says_whether_a_shell_asks_first(self): # The page shows "waiting for the terminal" only for shells that really ask with tempfile.NamedTemporaryFile('w', suffix='.json', delete=False) as f: json.dump({'servers': [{'id': 'ro', 'shell': {'allow': ['ls'], 'confirm': False}}, {'id': 'git', 'shell': {'allow': ['git']}}, {'id': 'tools', 'cli': [{'name': 'wc', 'command': ['wc']}]}]}, f) self.addCleanup(pathlib.Path(f.name).unlink) listed = {s['id']: s for s in mcp_bridge.Bridge(f.name, 0).server_list()} self.assertEqual((listed['ro']['kind'], listed['ro']['confirm']), ('shell', False)) self.assertTrue(listed['git']['confirm']) self.assertNotIn('confirm', listed['tools']) def test_startup_lines_name_the_config_and_each_shells_confirm_setting(self): with tempfile.NamedTemporaryFile('w', suffix='.json', delete=False) as f: json.dump({'servers': [{'id': 'ro', 'shell': {'allow': ['ls'], 'confirm': False}}, {'id': 'git', 'shell': {'allow': ['git']}}]}, f) self.addCleanup(pathlib.Path(f.name).unlink) lines = mcp_bridge.Bridge(f.name, 0).startup_lines() self.assertIn(str(pathlib.Path(f.name).resolve()), lines[0]) self.assertTrue(any('ro' in l and '免确认' in l for l in lines[1:]), lines) self.assertTrue(any('git' in l and '终端确认' in l for l in lines[1:]), lines) def test_prompt_names_the_config_file(self): with tempfile.NamedTemporaryFile('w', suffix='.json', delete=False) as f: json.dump({'servers': [{'id': 'git', 'shell': {'allow': ['echo']}}]}, f) self.addCleanup(pathlib.Path(f.name).unlink) bridge = mcp_bridge.Bridge(f.name, 0, confirm=FakeConfirm(False)) call(bridge.session('git'), 'echo hi') self.assertIn(str(pathlib.Path(f.name).resolve()), bridge.confirm.prompts[0]) def test_invalid_shell_configs_are_rejected(self): for bad in ({'allow': []}, {'allow': 'ls'}, {'allow': ['ls -la']}, {'allow': ['']}, {'allow': ['ls'], 'confirm': 'no'}, {}, {'allow': ['ls'], 'extra': 1}): with self.assertRaises(ValueError, msg=bad): self.load([{'id': 'sh', 'shell': bad}]) with self.assertRaises(ValueError): self.load([{'id': 'sh', 'shell': {'allow': ['ls']}, 'command': ['x']}]) class TerminalConfirmTest(unittest.TestCase): """The real prompt, on a pseudo-terminal standing in for the user's terminal.""" def setUp(self): self.master, slave = os.openpty() self.tty = os.ttyname(slave) os.close(slave) self.addCleanup(os.close, self.master) def answer_later(self, keys, delay=0.2): def type_keys(): time.sleep(delay) os.write(self.master, keys) threading.Thread(target=type_keys, daemon=True).start() def read_screen(self): out = b'' while True: try: import select if not select.select([self.master], [], [], 0.1)[0]: return out.decode('utf-8', 'replace') out += os.read(self.master, 4096) except OSError: return out.decode('utf-8', 'replace') def test_yes_approves_and_the_question_is_shown(self): self.answer_later(b'y\n') ok, why = mcp_shell.TerminalConfirm(self.tty, timeout=5)('运行 wc -l？') self.assertTrue(ok, why) self.assertIn('运行 wc -l？', self.read_screen()) def test_anything_else_refuses(self): for keys in (b'n\n', b'\n', b'yes please\n', b'Y E S\n'): self.answer_later(keys, 0.1) ok, why = mcp_shell.TerminalConfirm(self.tty, timeout=5)('?') self.assertFalse(ok, keys) self.assertIn('拒绝', why) def test_keys_typed_before_the_question_are_ignored(self): os.write(self.master, b'y\n') # typed ahead, before the prompt appears time.sleep(0.1) ok, _ = mcp_shell.TerminalConfirm(self.tty, timeout=0.5)('?') self.assertFalse(ok) def test_no_answer_times_out_as_refusal(self): ok, why = mcp_shell.TerminalConfirm(self.tty, timeout=0.3)('?') self.assertFalse(ok) self.assertIn('秒', why) def test_a_second_request_while_one_waits_is_refused_at_once(self): confirm = mcp_shell.TerminalConfirm(self.tty, timeout=1.5) first = threading.Thread(target=confirm, args=('first?',)) first.start() time.sleep(0.2) t0 = time.monotonic() ok, why = confirm('second?') self.assertFalse(ok) self.assertIn('正在等待', why) self.assertLess(time.monotonic() - t0, 0.5) first.join() def test_no_terminal_refuses(self): ok, why = mcp_shell.TerminalConfirm('/nonexistent/tty', timeout=1)('?') self.assertFalse(ok) self.assertIn('终端', why) if __name__ == '__main__': unittest.main()