diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..1e9b826 --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1 @@ +# xiaoli-cli test suite diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..c47a4ea --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,183 @@ +"""共享测试 fixtures""" +import os +import sys +import json +import tempfile +import shutil +import pytest + +# 确保项目根目录在 sys.path 中 +PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, PROJECT_ROOT) + + +@pytest.fixture +def tmp_dir(): + """创建临时目录,测试后清理""" + d = tempfile.mkdtemp(prefix='xiaoli_test_') + yield d + shutil.rmtree(d, ignore_errors=True) + + +@pytest.fixture +def sample_python_code(): + """安全的示例 Python 代码""" + return ''' +import json +import math + +result = math.sqrt(144) +data = {"value": result, "squared": result ** 2} +print(json.dumps(data)) +''' + + +@pytest.fixture +def dangerous_codes(): + """各种危险代码样本""" + return { + 'os_import': 'import os\nos.system("rm -rf /")', + 'subprocess': 'import subprocess\nsubprocess.run(["ls"])', + 'exec_call': 'exec("import os")', + 'eval_call': 'eval("__import__(\'os\')")', + 'dunder': 'x = (1).__class__.__bases__[0].__subclasses__()', + 'open_file': 'f = open("/etc/passwd")', + 'getattr_bypass': 'getattr(__builtins__, "exec")', + 'socket': 'import socket\ns = socket.socket()', + 'shutil': 'import shutil\nshutil.rmtree("/")', + 'compile': 'compile("import os", "", "exec")', + 'globals': 'globals()["__builtins__"]', + 'breakpoint': 'breakpoint()', + 'del_stmt': 'del some_var', + 'global_stmt': 'global x', + } + + +@pytest.fixture +def safe_codes(): + """各种安全代码样本""" + return { + 'math': 'import math\nresult = math.pi', + 'json': 'import json\nresult = json.dumps({"a": 1})', + 're': 'import re\nresult = re.findall(r"\\d+", "abc123")', + 'datetime': 'from datetime import datetime\nresult = datetime.now().year', + 'collections': 'from collections import Counter\nresult = Counter("aabbc")', + 'list_comp': 'result = [x**2 for x in range(10)]', + 'function': 'def add(a, b): return a + b\nresult = add(3, 4)', + 'string_ops': 'result = "hello world".upper().split()', + 'sorting': 'result = sorted([3,1,2], reverse=True)', + 'fizzbuzz': 'result = ["fizz"*(i%3==0)+"buzz"*(i%5==0) or str(i) for i in range(1,16)]', + } + + +@pytest.fixture +def config_file(tmp_dir): + """创建临时配置文件""" + config = { + "system": { + "default_engine": "ollama", + "max_history": 100, + "code_execution_timeout": 10, + }, + "api": { + "engines": { + "ollama": {"base_url": "http://localhost:11434"}, + "mimo": {"api_key": "test-key", "base_url": "http://test.com"}, + } + }, + "plugins": {}, + } + path = os.path.join(tmp_dir, "config.json") + with open(path, 'w') as f: + json.dump(config, f) + return path + + +@pytest.fixture +def history_dir(tmp_dir): + """创建临时历史目录""" + d = os.path.join(tmp_dir, "chat_history") + os.makedirs(d) + return d + + +@pytest.fixture +def plugins_dir(tmp_dir): + """创建带示例插件的目录""" + d = os.path.join(tmp_dir, "plugins") + os.makedirs(d) + + # __init__.py + with open(os.path.join(d, "__init__.py"), 'w') as f: + f.write("") + + # 简单插件 + with open(os.path.join(d, "test_plugin.py"), 'w') as f: + f.write(''' +class Plugin: + def __init__(self): + self.cli = None + + def set_cli(self, cli): + self.cli = cli + + def get_tool_info(self): + return { + "name": "test_tool", + "description": "A test tool", + "keywords": ["test"], + } + + def get_mcp_definition(self): + return { + "name": "test_tool", + "description": "A test tool", + "inputSchema": { + "type": "object", + "properties": { + "operation": {"type": "string"}, + }, + "required": ["operation"] + } + } + + def convert_mcp_args(self, arguments): + return arguments.get("operation", "") + + def handle(self, args): + return f"test result: {args}" +''') + + return d + + +@pytest.fixture +def skills_dir(tmp_dir): + """创建带示例 skill 的目录""" + d = os.path.join(tmp_dir, "skills") + os.makedirs(d) + + with open(os.path.join(d, "__init__.py"), 'w') as f: + f.write("") + + with open(os.path.join(d, "test_skill.py"), 'w') as f: + f.write(''' +class Skill: + def __init__(self): + self.cli = None + + def set_cli(self, cli): + self.cli = cli + + def get_tool_info(self): + return { + "name": "test_skill", + "description": "A test skill", + "keywords": ["test", "skill"], + } + + def handle(self, args): + return f"skill result: {args}" +''') + + return d diff --git a/tests/test_config.py b/tests/test_config.py new file mode 100644 index 0000000..df0a421 --- /dev/null +++ b/tests/test_config.py @@ -0,0 +1,137 @@ +"""配置管理测试""" +import os +import sys +import json +import pytest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from xiaoli import config as cfg + + +class TestConfig: + """配置加载/保存测试""" + + def setup_method(self): + """重置缓存""" + cfg._cache = None + + def test_default_config(self): + """无配置文件时使用默认值""" + # 临时指向不存在的文件 + original = cfg.CONFIG_FILE + cfg.CONFIG_FILE = '/tmp/nonexistent_config.json' + cfg._cache = None + + result = cfg.load() + assert 'system' in result + assert result['system']['default_engine'] == 'ollama' + + cfg.CONFIG_FILE = original + cfg._cache = None + + def test_load_existing_config(self, config_file): + """加载现有配置文件""" + original = cfg.CONFIG_FILE + cfg.CONFIG_FILE = config_file + cfg._cache = None + + result = cfg.load() + assert result['system']['default_engine'] == 'ollama' + assert result['system']['max_history'] == 100 + assert 'mimo' in result['api']['engines'] + + cfg.CONFIG_FILE = original + cfg._cache = None + + def test_get_specific_key(self, config_file): + """获取特定配置项""" + original = cfg.CONFIG_FILE + cfg.CONFIG_FILE = config_file + cfg._cache = None + + assert cfg.get('default_engine') == 'ollama' + assert cfg.get('max_history') == 100 + assert cfg.get('nonexistent', 'fallback') == 'fallback' + + cfg.CONFIG_FILE = original + cfg._cache = None + + def test_set_and_save(self, tmp_dir): + """设置并保存配置""" + config_path = os.path.join(tmp_dir, "config.json") + original = cfg.CONFIG_FILE + cfg.CONFIG_FILE = config_path + cfg._cache = None + + cfg.set('default_engine', 'mimo') + assert cfg.get('default_engine') == 'mimo' + + # 验证文件持久化 + with open(config_path, 'r') as f: + saved = json.load(f) + assert saved['system']['default_engine'] == 'mimo' + + cfg.CONFIG_FILE = original + cfg._cache = None + + def test_reload(self, config_file): + """强制重新加载""" + original = cfg.CONFIG_FILE + cfg.CONFIG_FILE = config_file + cfg._cache = None + + cfg.load() + assert cfg._cache is not None + + reloaded = cfg.reload() + assert cfg._cache is not None + assert 'system' in reloaded + + cfg.CONFIG_FILE = original + cfg._cache = None + + def test_caching(self, config_file): + """配置缓存机制""" + original = cfg.CONFIG_FILE + cfg.CONFIG_FILE = config_file + cfg._cache = None + + r1 = cfg.load() + r2 = cfg.load() + assert r1 is r2 # 同一对象(缓存) + + cfg.CONFIG_FILE = original + cfg._cache = None + + def test_corrupted_json(self, tmp_dir): + """损坏的 JSON 文件处理""" + config_path = os.path.join(tmp_dir, "config.json") + with open(config_path, 'w') as f: + f.write("{invalid json") + + original = cfg.CONFIG_FILE + cfg.CONFIG_FILE = config_path + cfg._cache = None + + result = cfg.load() + assert 'system' in result # 应回退到默认值 + + cfg.CONFIG_FILE = original + cfg._cache = None + + def test_missing_system_section(self, tmp_dir): + """缺少 system 节时的行为""" + config_path = os.path.join(tmp_dir, "config.json") + with open(config_path, 'w') as f: + json.dump({"api": {}}, f) + + original = cfg.CONFIG_FILE + cfg.CONFIG_FILE = config_path + cfg._cache = None + + assert cfg.get('default_engine') is None + assert cfg.get('default_engine', 'ollama') == 'ollama' + + cfg.CONFIG_FILE = original + cfg._cache = None diff --git a/tests/test_conversation.py b/tests/test_conversation.py new file mode 100644 index 0000000..a978785 --- /dev/null +++ b/tests/test_conversation.py @@ -0,0 +1,226 @@ +"""对话响应解析器测试""" +import os +import sys +import pytest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from xiaoli.conversation import parse_response, process_thinking, process_code_blocks, extract_tool_calls + + +# ═══════════════════════════════════════════════ +# 1. parse_response 测试 +# ═══════════════════════════════════════════════ + +class TestParseResponse: + """响应解析测试""" + + def test_pure_text(self): + text, data = parse_response("hello world") + assert text == "hello world" + assert data is None + + def test_empty_response(self): + text, data = parse_response("") + assert text == "" + assert data is None + + def test_pure_json_dict(self): + resp = '{"action": "use_tool", "tool": "code_editor", "args": "read foo.py"}' + text, data = parse_response(resp) + assert text == "" + assert isinstance(data, dict) + assert data['action'] == 'use_tool' + assert data['tool'] == 'code_editor' + + def test_pure_json_list(self): + resp = '[{"action": "use_tool", "tool": "t1", "args": "a1"}]' + text, data = parse_response(resp) + assert text == "" + assert isinstance(data, list) + + def test_json_in_code_block(self): + resp = 'Here is the command:\n```json\n{"action": "use_tool", "tool": "git_tools", "args": "status"}\n```' + text, data = parse_response(resp) + assert 'command' in text.lower() or 'Here' in text + assert isinstance(data, dict) + assert data['tool'] == 'git_tools' + + def test_text_before_json(self): + resp = 'I will edit the file.\n{"action": "use_tool", "tool": "code_editor", "args": "edit foo.py"}' + text, data = parse_response(resp) + assert 'edit' in text.lower() + assert isinstance(data, dict) + assert data['action'] == 'use_tool' + + def test_continue_action(self): + resp = '{"action": "continue", "content": "more info"}' + text, data = parse_response(resp) + assert text == "" + assert isinstance(data, dict) + assert data['action'] == 'continue' + + def test_inline_json_use_tool(self): + resp = 'Let me check. {"action":"use_tool","tool":"code_search","args":"find . test"}' + text, data = parse_response(resp) + assert isinstance(data, dict) + assert data['tool'] == 'code_search' + + def test_inline_json_continue(self): + resp = 'Some text. {"action":"continue","content":"extra"}' + text, data = parse_response(resp) + assert isinstance(data, dict) + assert data['action'] == 'continue' + + def test_multiple_json_per_line(self): + resp = '{"action":"use_tool","tool":"t1","args":"a1"}\n{"action":"use_tool","tool":"t2","args":"a2"}' + text, data = parse_response(resp) + # Should pick up at least one + assert data is not None + + def test_invalid_json_ignored(self): + resp = 'This is not json {broken' + text, data = parse_response(resp) + assert text == resp + assert data is None + + def test_tool_call_tag_format(self): + resp = '{"action":"use_tool","tool":"test","args":"arg"}' + text, data = parse_response(resp) + assert isinstance(data, dict) + assert data['tool'] == 'test' + + def test_message_field(self): + resp = '{"message": "done processing"}' + text, data = parse_response(resp) + assert isinstance(data, dict) + assert data['message'] == 'done processing' + + +# ═══════════════════════════════════════════════ +# 2. process_thinking 测试 +# ═══════════════════════════════════════════════ + +class TestProcessThinking: + """思考标记处理测试""" + + def test_think_tag(self): + resp = 'This is the reply' + result = process_thinking(resp) + assert result == 'This is the reply' + + def test_think_tag_with_thinking(self): + resp = 'I am thinking deeplyHere is my answer' + result = process_thinking(resp) + assert 'answer' in result + + def test_no_think_tag(self): + resp = 'Just a normal response' + result = process_thinking(resp) + assert result == resp + + def test_empty_after_think(self): + resp = '' + result = process_thinking(resp) + assert result == '' or result is not None + + def test_custom_markers(self): + class FakeEngine: + thinking_start_marker = '' + thinking_end_marker = '' + + resp = 'deep thoughtThe answer is 42' + result = process_thinking(resp, FakeEngine()) + assert '42' in result + + +# ═══════════════════════════════════════════════ +# 3. process_code_blocks 测试 +# ═══════════════════════════════════════════════ + +class TestProcessCodeBlocks: + """代码块处理测试""" + + def test_triple_quote_block(self): + text = 'before\n"""\ncode here\n"""\nafter' + result = process_code_blocks(text) + assert 'before' in result + assert 'after' in result + + def test_no_code_blocks(self): + text = 'plain text no blocks' + result = process_code_blocks(text) + assert result == text + + +# ═══════════════════════════════════════════════ +# 4. extract_tool_calls 测试 +# ═══════════════════════════════════════════════ + +class TestExtractToolCalls: + """工具调用提取测试""" + + def test_single_use_tool(self): + data = {"action": "use_tool", "tool": "code_editor", "args": "edit foo.py"} + calls = extract_tool_calls(data) + assert len(calls) == 1 + assert calls[0]['tool'] == 'code_editor' + + def test_list_of_use_tool(self): + data = [ + {"action": "use_tool", "tool": "t1", "args": "a1"}, + {"action": "use_tool", "tool": "t2", "args": "a2"}, + ] + calls = extract_tool_calls(data) + assert len(calls) == 2 + + def test_simplified_format(self): + data = {"tool": "git_tools", "args": "status"} + calls = extract_tool_calls(data) + assert len(calls) == 1 + assert calls[0]['tool'] == 'git_tools' + + def test_openai_function_call(self): + data = {"function_call": {"name": "code_editor", "arguments": '{"operation": "read"}'}} + calls = extract_tool_calls(data) + assert len(calls) == 1 + assert calls[0]['tool'] == 'code_editor' + + def test_mcp_format(self): + data = {"name": "code_search", "arguments": {"operation": "find", "directory": "."}} + calls = extract_tool_calls(data) + assert len(calls) == 1 + assert calls[0]['tool'] == 'code_search' + + def test_empty_data(self): + assert extract_tool_calls([]) == [] + assert extract_tool_calls({}) == [] + assert extract_tool_calls(None) == [] + assert extract_tool_calls("string") == [] + + def test_non_tool_dict(self): + data = {"message": "just text", "continue": True} + calls = extract_tool_calls(data) + # continue action should not be extracted as tool call + assert len(calls) == 0 + + def test_mixed_list(self): + data = [ + {"action": "use_tool", "tool": "t1", "args": "a1"}, + {"action": "continue", "content": "more"}, + ] + calls = extract_tool_calls(data) + assert len(calls) == 1 + assert calls[0]['tool'] == 't1' + + def test_function_call_string_args(self): + data = {"function_call": {"name": "test", "arguments": '{"key": "val"}'}} + calls = extract_tool_calls(data) + assert len(calls) == 1 + assert calls[0]['arguments'] == {"key": "val"} + + def test_function_call_dict_args(self): + data = {"function_call": {"name": "test", "arguments": {"key": "val"}}} + calls = extract_tool_calls(data) + assert len(calls) == 1 + assert calls[0]['arguments'] == {"key": "val"} diff --git a/tests/test_engines.py b/tests/test_engines.py new file mode 100644 index 0000000..f9c7a02 --- /dev/null +++ b/tests/test_engines.py @@ -0,0 +1,107 @@ +"""AI 引擎管理器测试""" +import os +import sys +import pytest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from xiaoli.engines import EngineManager + + +class TestEngineManager: + """EngineManager 测试""" + + def test_init_empty(self): + mgr = EngineManager() + assert mgr.engines == {} + assert mgr.current is None + + def test_set_default_existing(self): + mgr = EngineManager() + mgr.engines = {'test': 'fake_engine'} + assert mgr.set_default('test') + assert mgr.current == 'fake_engine' + + def test_set_default_missing(self): + mgr = EngineManager() + assert not mgr.set_default('nonexistent') + + def test_switch_existing(self): + mgr = EngineManager() + mgr.engines = {'a': 'engine_a', 'b': 'engine_b'} + assert mgr.switch('b') + assert mgr.current == 'engine_b' + + def test_switch_missing(self): + mgr = EngineManager() + assert not mgr.switch('nope') + + def test_current_name_none(self): + mgr = EngineManager() + assert mgr.current_name() == '未设置' + + def test_current_name_set(self): + mgr = EngineManager() + mgr.current = type('E', (), {'name': 'ollama'})() + assert mgr.current_name() == 'ollama' + + def test_names(self): + mgr = EngineManager() + mgr.engines = {'a': 1, 'b': 2, 'c': 3} + names = mgr.names() + assert set(names) == {'a', 'b', 'c'} + + def test_generate_no_engine(self): + mgr = EngineManager() + result = mgr.generate("hello") + assert '错误' in result or 'error' in result.lower() + + def test_generate_with_engine(self): + mgr = EngineManager() + + class FakeEngine: + name = 'fake' + def generate_response(self, user_input, system_prompt=None): + return f"echo: {user_input}" + + mgr.current = FakeEngine() + result = mgr.generate("test input") + assert 'test input' in result + + def test_load_all_missing_dir(self): + mgr = EngineManager() + result = mgr.load_all('/tmp/nonexistent_engines_dir_12345') + assert result == {} + + def test_find_class_patterns(self): + mgr = EngineManager() + import types + mod = types.ModuleType('test') + + class TestAI: + def generate_response(self, **kw): + pass + mod.TestAI = TestAI + + found = mgr._find_class(mod, 'test') + assert found == 'TestAI' + + def test_find_class_no_match(self): + mgr = EngineManager() + import types + mod = types.ModuleType('test') + mod.Foo = 42 # not a class ending in AI + found = mgr._find_class(mod, 'test') + assert found is None + + def test_load_all_real_engines(self): + """测试加载实际引擎目录(不验证引擎可用性,只验证不崩溃)""" + mgr = EngineManager() + engines_dir = os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + 'ai_engines' + ) + if os.path.exists(engines_dir): + mgr.load_all(engines_dir) + # 至少应该尝试加载(可能因缺少依赖而失败) + # 但不应崩溃 diff --git a/tests/test_history.py b/tests/test_history.py new file mode 100644 index 0000000..028162c --- /dev/null +++ b/tests/test_history.py @@ -0,0 +1,114 @@ +"""对话历史管理测试""" +import os +import sys +import json +import pytest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from xiaoli.conversation.history import History + + +class TestHistory: + """History 类测试""" + + def test_add_message(self, history_dir): + h = History(history_dir) + h.add("user", "hello") + h.add("assistant", "hi there") + assert len(h.messages) == 2 + assert h.messages[0] == {"role": "user", "content": "hello"} + assert h.messages[1] == {"role": "assistant", "content": "hi there"} + + def test_get_messages(self, history_dir): + h = History(history_dir) + h.add("user", "q1") + h.add("assistant", "a1") + msgs = h.get() + assert len(msgs) == 2 + assert msgs[0]['role'] == 'user' + + def test_clear(self, history_dir): + h = History(history_dir) + h.add("user", "hello") + h.add("assistant", "hi") + h.clear() + assert len(h.messages) == 0 + + def test_max_items(self, history_dir): + h = History(history_dir, max_items=3) + for i in range(5): + h.add("user", f"msg {i}") + assert len(h.messages) == 3 + assert h.messages[0]['content'] == 'msg 2' + + def test_save_and_load(self, history_dir): + h = History(history_dir) + h.add("user", "hello") + h.add("assistant", "hi") + assert h.save("test_chat") + + h2 = History(history_dir) + assert h2.load("test_chat") + assert len(h2.messages) == 2 + assert h2.messages[0]['content'] == 'hello' + + def test_load_nonexistent(self, history_dir, capsys): + h = History(history_dir) + result = h.load("nonexistent_chat") + assert not result + + def test_list_saved(self, history_dir, capsys): + h = History(history_dir) + h.add("user", "test") + h.save("chat1") + h.save("chat2") + + h2 = History(history_dir) + h2.list_saved() # should not raise + captured = capsys.readouterr() + assert 'chat1' in captured.out or 'chat2' in captured.out + + def test_list_saved_empty(self, history_dir, capsys): + h = History(history_dir) + h.list_saved() + captured = capsys.readouterr() + assert 'no' in captured.out.lower() or '没有' in captured.out + + def test_set_all(self, history_dir): + h = History(history_dir) + msgs = [ + {"role": "user", "content": "q"}, + {"role": "assistant", "content": "a"}, + ] + h.set_all(msgs) + assert len(h.messages) == 2 + assert h.messages == msgs + + def test_save_creates_directory(self, tmp_dir): + nested_dir = os.path.join(tmp_dir, "deep", "nested", "history") + h = History(nested_dir) + h.add("user", "test") + h.save("test") + assert os.path.exists(os.path.join(nested_dir, "test.json")) + + def test_save_file_format(self, history_dir): + h = History(history_dir) + h.add("user", "hello") + h.save("format_test") + + path = os.path.join(history_dir, "format_test.json") + with open(path, 'r') as f: + data = json.load(f) + assert 'timestamp' in data + assert 'conversation' in data + assert len(data['conversation']) == 1 + + def test_load_corrupted_file(self, history_dir): + path = os.path.join(history_dir, "corrupt.json") + with open(path, 'w') as f: + f.write("{bad json") + + h = History(history_dir) + result = h.load("corrupt") + assert not result diff --git a/tests/test_integration.py b/tests/test_integration.py new file mode 100644 index 0000000..b30edc0 --- /dev/null +++ b/tests/test_integration.py @@ -0,0 +1,154 @@ +"""集成测试 - 模块间协作""" +import os +import sys +import json +import pytest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from xiaoli.sandbox import Sandbox +from xiaoli.conversation import parse_response, extract_tool_calls +from xiaoli.conversation.history import History + + +class TestSandboxConversationIntegration: + """沙箱 + 对话解析 集成测试""" + + def test_sandbox_result_parsed_as_tool_output(self): + """沙箱执行结果可以正确传递给对话系统""" + sandbox = Sandbox(timeout=10, mode='threaded') + result = sandbox.run('result = [x**2 for x in range(5)]') + assert result['success'] + + # 模拟工具结果注入对话 + tool_msg = f"工具执行结果:\n{result['result']}" + text, data = parse_response(tool_msg) + assert data is None # 纯文本,无工具调用 + assert '0' in text and '16' in text + + sandbox.cleanup() + + def test_sandbox_failure_triggers_error_handling(self): + """沙箱失败时应正确处理""" + sandbox = Sandbox(timeout=10, mode='threaded') + result = sandbox.run('import os') # should fail + assert not result['success'] + + # 错误信息应可解析 + assert result['result'] + assert len(result['result']) > 0 + + sandbox.cleanup() + + +class TestHistoryConversationIntegration: + """历史 + 对话解析 集成测试""" + + def test_full_conversation_roundtrip(self, history_dir): + """完整的对话存取流程""" + h = History(history_dir) + + # 模拟对话 + h.add("user", "帮我写一个排序函数") + h.add("assistant", "好的,这是一个快速排序:\n```python\ndef quicksort(arr):\n ...\n```") + h.add("user", "能改成归并排序吗?") + + # 保存 + assert h.save("sorting_chat") + + # 加载到新实例 + h2 = History(history_dir) + assert h2.load("sorting_chat") + assert len(h2.messages) == 3 + assert "quicksort" in h2.messages[1]['content'] + + def test_conversation_with_tool_calls(self, history_dir): + """带工具调用的对话历史""" + h = History(history_dir) + + # 用户请求 + h.add("user", "查看 main.py 的第 10-20 行") + + # AI 的工具调用 + tool_call = '{"action":"use_tool","tool":"code_editor","args":"read_range main.py 10 20"}' + text, data = parse_response(tool_call) + assert data is not None + calls = extract_tool_calls(data) + assert len(calls) == 1 + assert calls[0]['tool'] == 'code_editor' + + # 记录工具结果 + h.add("assistant", f"工具调用: {json.dumps(data)}") + h.add("assistant", "工具执行结果:\n line 10: def main():") + + assert len(h.messages) == 3 + + +class TestSandboxSecurityIntegration: + """沙箱安全 + 工具链 集成测试""" + + def test_safe_code_executes_in_chain(self): + """安全代码在工具链中正常执行""" + sandbox = Sandbox(timeout=10, mode='threaded') + + # 步骤 1: 数学计算 + r1 = sandbox.run('import math\nresult = math.factorial(10)') + assert r1['success'] + assert '3628800' in r1['result'] + + # 步骤 2: 使用上一步结果 + r2 = sandbox.run( + 'result = f"factorial(10) = {val}"', + variables={'val': r1['result']} + ) + assert r2['success'] + assert '3628800' in r2['result'] + + sandbox.cleanup() + + def test_blocked_code_stops_chain(self): + """被阻止的代码应中断工具链""" + sandbox = Sandbox(timeout=10, mode='threaded') + + r = sandbox.run('import os\nresult = os.listdir(".")') + assert not r['success'] + # 后续步骤不应执行 + assert r['result'] + + sandbox.cleanup() + + def test_subprocess_isolation(self): + """子进程模式下文件系统隔离""" + sandbox = Sandbox(timeout=10, mode='subprocess') + result = sandbox.run( + 'import tempfile\nresult = tempfile.gettempdir()' + ) + # 子进程应使用沙箱临时目录 + if result['success']: + assert 'xiaoli_sandbox' in result['result'] or result['success'] + + sandbox.cleanup() + + +class TestPromptIntegration: + """系统提示词集成测试""" + + def test_prompt_builder(self): + from xiaoli import prompt as prompt_builder + + # 无工具 + p1 = prompt_builder.build() + assert '小狸' in p1 + assert '编程' in p1 or 'code' in p1.lower() + + # 带工具 + p2 = prompt_builder.build(tool_prompts="- code_editor: edit code") + assert 'code_editor' in p2 + + # 带 tool_search + p3 = prompt_builder.build(has_tool_search=True) + assert 'tool_search' in p3 + + # Clawli 模式 + p4 = prompt_builder.build(is_clawli=True) + assert 'Clawli' in p4 or 'clawli' in p4 or '远程' in p4 diff --git a/tests/test_plugins.py b/tests/test_plugins.py new file mode 100644 index 0000000..f26a60b --- /dev/null +++ b/tests/test_plugins.py @@ -0,0 +1,240 @@ +"""插件系统测试""" +import os +import sys +import pytest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + + +# ═══════════════════════════════════════════════ +# 1. code_editor 插件测试 +# ═══════════════════════════════════════════════ + +class TestCodeEditor: + """代码编辑器插件测试""" + + @pytest.fixture(autouse=True) + def setup(self, tmp_dir): + sys.path.insert(0, os.path.join(PROJECT_ROOT, 'plugins')) + from code_editor import Plugin + self.plugin = Plugin() + self.tmp_dir = tmp_dir + + def test_tool_info(self): + info = self.plugin.get_tool_info() + assert info['name'] == 'code_editor' + assert 'description' in info + + def test_mcp_definition(self): + mcp = self.plugin.get_mcp_definition() + assert mcp['name'] == 'code_editor' + assert 'inputSchema' in mcp + assert 'operation' in mcp['inputSchema']['properties'] + + def test_convert_mcp_args(self): + args = self.plugin.convert_mcp_args({"operation": "edit", "path": "test.py"}) + assert 'edit' in args + + def test_create_file(self): + filepath = os.path.join(self.tmp_dir, 'test.py') + result = self.plugin.handle(f'create {filepath}\nprint("hello")') + assert os.path.exists(filepath) + with open(filepath) as f: + assert 'print("hello")' in f.read() + + def test_read_range(self): + filepath = os.path.join(self.tmp_dir, 'read_test.py') + with open(filepath, 'w') as f: + f.write('line1\nline2\nline3\nline4\nline5\n') + result = self.plugin.handle(f'read_range {filepath} 2 4') + assert 'line2' in result + assert 'line3' in result + + def test_edit_file(self): + filepath = os.path.join(self.tmp_dir, 'edit_test.py') + with open(filepath, 'w') as f: + f.write('def old():\n pass\n') + result = self.plugin.handle(f'edit {filepath} def old():\n pass <<<>>> def new():\n return True') + with open(filepath) as f: + content = f.read() + assert 'def new():' in content + assert 'return True' in content + + def test_diff_files(self): + f1 = os.path.join(self.tmp_dir, 'a.py') + f2 = os.path.join(self.tmp_dir, 'b.py') + with open(f1, 'w') as f: + f.write('hello\nworld\n') + with open(f2, 'w') as f: + f.write('hello\nuniverse\n') + result = self.plugin.handle(f'diff {f1} {f2}') + assert result is not None # should produce diff output + + def test_search_code(self): + filepath = os.path.join(self.tmp_dir, 'search.py') + with open(filepath, 'w') as f: + f.write('def foo():\n return 42\ndef bar():\n return foo()\n') + result = self.plugin.handle(f'search {self.tmp_dir} main *.py') + assert 'main' in result + + def test_write_file(self): + filepath = os.path.join(self.tmp_dir, 'write.py') + result = self.plugin.handle(f'write {filepath}\nresult = 42\n') + assert os.path.exists(filepath) + with open(filepath) as f: + assert 'result = 42' in f.read() + + def test_append_file(self): + filepath = os.path.join(self.tmp_dir, 'append.py') + with open(filepath, 'w') as f: + f.write('line1\n') + self.plugin.handle(f'append {filepath}\nline2\n') + with open(filepath) as f: + content = f.read() + assert 'line1' in content + assert 'line2' in content + + +# ═══════════════════════════════════════════════ +# 2. code_search 插件测试 +# ═══════════════════════════════════════════════ + +class TestCodeSearch: + """代码搜索插件测试""" + + @pytest.fixture(autouse=True) + def setup(self, tmp_dir): + sys.path.insert(0, os.path.join(PROJECT_ROOT, 'plugins')) + from code_search import Plugin + self.plugin = Plugin() + self.tmp_dir = tmp_dir + + # 创建测试文件 + with open(os.path.join(tmp_dir, 'main.py'), 'w') as f: + f.write('import os\ndef main():\n print("hello")\n\nif __name__ == "__main__":\n main()\n') + with open(os.path.join(tmp_dir, 'utils.py'), 'w') as f: + f.write('def helper():\n return 42\ndef process(data):\n return [x*2 for x in data]\n') + with open(os.path.join(tmp_dir, 'test_main.py'), 'w') as f: + f.write('# TODO: add tests\ndef test_main():\n assert True\n') + + def test_tool_info(self): + info = self.plugin.get_tool_info() + assert info['name'] == 'code_search' + + def test_structure(self): + result = self.plugin.handle(f'structure {self.tmp_dir}') + assert 'main.py' in result + assert 'utils.py' in result + + def test_find(self): + result = self.plugin.handle(f'find {self.tmp_dir} def *.py') + assert 'main' in result or 'helper' in result or 'def' in result + + def test_symbols(self): + result = self.plugin.handle(f'symbols {self.tmp_dir} main.py') + assert 'main' in result + + def test_todo(self): + result = self.plugin.handle(f'todo {self.tmp_dir}') + assert 'TODO' in result + + def test_stats(self): + result = self.plugin.handle(f'stats {self.tmp_dir}') + assert result is not None + + def test_context(self): + filepath = os.path.join(self.tmp_dir, 'main.py') + result = self.plugin.handle(f'context {filepath} 2 2') + assert 'main' in result or 'def' in result + + +# ═══════════════════════════════════════════════ +# 3. git_tools 插件测试 +# ═══════════════════════════════════════════════ + +class TestGitTools: + """Git 工具插件测试""" + + @pytest.fixture(autouse=True) + def setup(self, tmp_dir): + sys.path.insert(0, os.path.join(PROJECT_ROOT, 'plugins')) + from git_tools import Plugin + self.plugin = Plugin() + self.tmp_dir = tmp_dir + + def test_tool_info(self): + info = self.plugin.get_tool_info() + assert info['name'] == 'git_tools' + + def test_mcp_definition(self): + mcp = self.plugin.get_mcp_definition() + assert mcp['name'] == 'git_tools' + assert 'inputSchema' in mcp + + def test_status_non_git_dir(self): + result = self.plugin.handle(f'status') + # 在非 git 目录应返回错误或提示 + assert result is not None + + def test_help_text(self): + assert 'git' in self.plugin.usage.lower() or 'Git' in self.plugin.usage + + +# ═══════════════════════════════════════════════ +# 4. unified_tool_manager 测试 +# ═══════════════════════════════════════════════ + +class TestUnifiedToolManager: + """统一工具管理器测试""" + + @pytest.fixture(autouse=True) + def setup(self, plugins_dir, skills_dir): + from unified_tool_manager import UnifiedToolManager + self.manager = UnifiedToolManager() + self.plugins_dir = plugins_dir + self.skills_dir = skills_dir + + def test_initialize(self): + self.manager.initialize(self.plugins_dir, self.skills_dir) + assert len(self.manager._tools) > 0 + + def test_tools_registered(self): + self.manager.initialize(self.plugins_dir, self.skills_dir) + assert 'test_tool' in self.manager._tools + + def test_protocol_map(self): + self.manager.initialize(self.plugins_dir, self.skills_dir) + assert 'test_tool' in self.manager._protocol_map + + def test_execute_plugin(self): + self.manager.initialize(self.plugins_dir, self.skills_dir) + result = self.manager.execute('test_tool', 'hello') + assert result is not None + assert 'hello' in str(result) or 'test' in str(result).lower() + + def test_execute_nonexistent(self): + self.manager.initialize(self.plugins_dir, self.skills_dir) + result = self.manager.execute('nonexistent_tool', 'args') + assert result is None or 'not found' in str(result).lower() or '未找到' in str(result) + + def test_mcp_cache(self): + self.manager.initialize(self.plugins_dir, self.skills_dir) + cache = self.manager._mcp_tools_cache + assert cache is not None + assert len(cache) > 0 + + def test_empty_dirs(self, tmp_dir): + empty_plugins = os.path.join(tmp_dir, 'empty_p') + empty_skills = os.path.join(tmp_dir, 'empty_s') + os.makedirs(empty_plugins) + os.makedirs(empty_skills) + with open(os.path.join(empty_plugins, '__init__.py'), 'w') as f: + pass + with open(os.path.join(empty_skills, '__init__.py'), 'w') as f: + pass + + mgr = __import__('unified_tool_manager').UnifiedToolManager() + mgr.initialize(empty_plugins, empty_skills) + assert len(mgr._tools) == 0 diff --git a/tests/test_sandbox.py b/tests/test_sandbox.py new file mode 100644 index 0000000..4d3e4d5 --- /dev/null +++ b/tests/test_sandbox.py @@ -0,0 +1,576 @@ +"""沙箱安全测试 - 测试所有安全层级""" +import os +import sys +import time +import pytest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from xiaoli.sandbox import Sandbox, SandboxResult, ASTSecurityChecker + + +# ═══════════════════════════════════════════════ +# 1. AST 安全检查器测试 +# ═══════════════════════════════════════════════ + +class TestASTSecurityChecker: + """AST 静态分析测试""" + + def setup_method(self): + self.checker = ASTSecurityChecker() + + # ── 应该拒绝的 ── + + def test_block_os_import(self): + safe, violations = self.checker.check('import os') + assert not safe + assert any('os' in v for v in violations) + + def test_block_subprocess_import(self): + safe, violations = self.checker.check('import subprocess') + assert not safe + + def test_block_from_import(self): + safe, violations = self.checker.check('from os import system') + assert not safe + assert any('os' in v for v in violations) + + def test_block_exec_call(self): + safe, violations = self.checker.check('exec("print(1)")') + assert not safe + assert any('exec' in v for v in violations) + + def test_block_eval_call(self): + safe, violations = self.checker.check('x = eval("1+1")') + assert not safe + + def test_block_compile(self): + safe, violations = self.checker.check('compile("x=1", "", "exec")') + assert not safe + + def test_block_globals(self): + safe, violations = self.checker.check('globals()') + assert not safe + + def test_block_locals(self): + safe, violations = self.checker.check('locals()') + assert not safe + + def test_block_getattr(self): + safe, violations = self.checker.check('getattr(obj, "name")') + assert not safe + + def test_block_setattr(self): + safe, violations = self.checker.check('setattr(obj, "name", 1)') + assert not safe + + def test_block_dunder_class(self): + safe, violations = self.checker.check('x = object.__class__') + assert not safe + + def test_block_dunder_subclasses(self): + safe, violations = self.checker.check('x = object.__subclasses__()') + assert not safe + + def test_block_open(self): + safe, violations = self.checker.check('f = open("/etc/passwd")') + assert not safe + + def test_block_breakpoint(self): + safe, violations = self.checker.check('breakpoint()') + assert not safe + + def test_block_socket_import(self): + safe, violations = self.checker.check('import socket') + assert not safe + + def test_block_pickle_import(self): + safe, violations = self.checker.check('import pickle') + assert not safe + + def test_block_ctypes_import(self): + safe, violations = self.checker.check('import ctypes') + assert not safe + + def test_block_importlib(self): + safe, violations = self.checker.check('import importlib') + assert not safe + + def test_block_del(self): + safe, violations = self.checker.check('del x') + assert not safe + + def test_block_global_stmt(self): + safe, violations = self.checker.check('global x') + assert not safe + + def test_block_nested_dunder(self): + safe, violations = self.checker.check( + 'x = (1).__class__.__bases__[0].__subclasses__()' + ) + assert not safe + + def test_block_sys_import(self): + safe, violations = self.checker.check('import sys') + assert not safe + + def test_block_shutil(self): + safe, violations = self.checker.check('import shutil') + assert not safe + + def test_block_urllib(self): + safe, violations = self.checker.check('import urllib') + assert not safe + + def test_block_http(self): + safe, violations = self.checker.check('import http') + assert not safe + + # ── 应该允许的 ── + + def test_allow_math(self): + safe, _ = self.checker.check('import math\nresult = math.pi') + assert safe + + def test_allow_json(self): + safe, _ = self.checker.check('import json\nresult = json.dumps({"a": 1})') + assert safe + + def test_allow_re(self): + safe, _ = self.checker.check('import re\nresult = re.findall(r"\\d+", "abc")') + assert safe + + def test_allow_datetime(self): + safe, _ = self.checker.check('from datetime import datetime\nresult = datetime.now()') + assert safe + + def test_allow_collections(self): + safe, _ = self.checker.check('from collections import Counter\nresult = Counter("aab")') + assert safe + + def test_allow_list_comprehension(self): + safe, _ = self.checker.check('result = [x**2 for x in range(10)]') + assert safe + + def test_allow_function_def(self): + safe, _ = self.checker.check('def add(a, b): return a + b\nresult = add(1, 2)') + assert safe + + def test_allow_string_ops(self): + safe, _ = self.checker.check('result = "hello".upper()') + assert safe + + def test_allow_sorting(self): + safe, _ = self.checker.check('result = sorted([3,1,2])') + assert safe + + def test_allow_typing(self): + safe, _ = self.checker.check('from typing import List\nresult: List[int] = [1,2,3]') + assert safe + + def test_allow_enum(self): + safe, _ = self.checker.check('from enum import Enum') + assert safe + + def test_allow_decimal(self): + safe, _ = self.checker.check('from decimal import Decimal\nresult = Decimal("3.14")') + assert safe + + def test_allow_itertools(self): + safe, _ = self.checker.check('import itertools\nresult = list(itertools.chain([1],[2]))') + assert safe + + # ── 语法错误 ── + + def test_syntax_error(self): + safe, violations = self.checker.check('def foo(') + assert not safe + assert any('syntax' in v.lower() for v in violations) + + def test_empty_code(self): + safe, _ = self.checker.check('') + assert safe + + # ── 多种违规组合 ── + + def test_multiple_violations(self): + code = 'import os\nimport socket\nexec("x=1")\nglobals()' + safe, violations = self.checker.check(code) + assert not safe + assert len(violations) >= 4 + + +# ═══════════════════════════════════════════════ +# 2. Sandbox 核心功能测试 +# ═══════════════════════════════════════════════ + +class TestSandboxBasic: + """沙箱基本功能测试""" + + def setup_method(self): + self.sandbox = Sandbox(timeout=10, mode='threaded') + + def teardown_method(self): + self.sandbox.cleanup() + + def test_simple_execution(self): + result = self.sandbox.run('result = 1 + 2') + assert result['success'] + assert result['result'] == '3' + + def test_string_manipulation(self): + result = self.sandbox.run('result = "hello world".split()') + assert result['success'] + assert 'hello' in result['result'] + + def test_math_operations(self): + result = self.sandbox.run('import math\nresult = math.sqrt(144)') + assert result['success'] + assert '12.0' in result['result'] + + def test_json_operations(self): + result = self.sandbox.run( + 'import json\nresult = json.dumps({"key": "value"})' + ) + assert result['success'] + assert 'key' in result['result'] + + def test_print_capture(self): + result = self.sandbox.run('print("hello")\nresult = "done"') + assert result['success'] + assert 'hello' in result['stdout'] + + def test_variables_injection(self): + result = self.sandbox.run( + 'result = x + y', + variables={'x': 10, 'y': 20} + ) + assert result['success'] + assert '30' in result['result'] + + def test_function_definition(self): + code = ''' +def factorial(n): + if n <= 1: + return 1 + return n * factorial(n - 1) +result = factorial(5) +''' + result = self.sandbox.run(code) + assert result['success'] + assert '120' in result['result'] + + def test_list_comprehension(self): + result = self.sandbox.run('result = [x**2 for x in range(5)]') + assert result['success'] + assert '[0, 1, 4, 9, 16]' in result['result'] + + def test_dict_operations(self): + code = ''' +data = {"a": 1, "b": 2, "c": 3} +result = {k: v*2 for k, v in data.items()} +''' + result = self.sandbox.run(code) + assert result['success'] + + def test_datetime(self): + code = 'from datetime import datetime\nresult = datetime.now().year' + result = self.sandbox.run(code) + assert result['success'] + + def test_collections_counter(self): + code = 'from collections import Counter\nresult = Counter("aabbc")' + result = self.sandbox.run(code) + assert result['success'] + + def test_regex(self): + code = 'import re\nresult = re.findall(r"\\d+", "abc123def456")' + result = self.sandbox.run(code) + assert result['success'] + assert '123' in result['result'] + + def test_no_result_variable(self): + result = self.sandbox.run('x = 42') + assert result['success'] + assert 'no return' in result['result'].lower() or 'ok' in result['result'].lower() + + def test_syntax_error(self): + result = self.sandbox.run('def foo(') + assert not result['success'] + + def test_runtime_error(self): + result = self.sandbox.run('result = 1 / 0') + assert not result['success'] + assert 'error' in result['result'].lower() or 'division' in result['result'].lower() + + +# ═══════════════════════════════════════════════ +# 3. 安全拒绝测试 +# ═══════════════════════════════════════════════ + +class TestSandboxSecurity: + """沙箱安全拒绝测试""" + + def setup_method(self): + self.sandbox = Sandbox(timeout=10, mode='threaded') + + def teardown_method(self): + self.sandbox.cleanup() + + def test_reject_os_import(self): + result = self.sandbox.run('import os\nresult = os.getcwd()') + assert not result['success'] + + def test_reject_subprocess(self): + result = self.sandbox.run('import subprocess\nsubprocess.run(["ls"])') + assert not result['success'] + + def test_reject_exec(self): + result = self.sandbox.run('exec("result = 1")') + assert not result['success'] + + def test_reject_eval(self): + result = self.sandbox.run('result = eval("1+1")') + assert not result['success'] + + def test_reject_open(self): + result = self.sandbox.run('f = open("/etc/passwd")\nresult = f.read()') + assert not result['success'] + + def test_reject_socket(self): + result = self.sandbox.run('import socket\ns = socket.socket()') + assert not result['success'] + + def test_reject_pickle(self): + result = self.sandbox.run('import pickle\nresult = pickle.dumps({})') + assert not result['success'] + + def test_reject_sys(self): + result = self.sandbox.run('import sys\nresult = sys.version') + assert not result['success'] + + def test_reject_getattr(self): + result = self.sandbox.run('result = getattr(str, "upper")') + assert not result['success'] + + def test_reject_dunder(self): + result = self.sandbox.run('result = object.__class__') + assert not result['success'] + + def test_reject_breakpoint(self): + result = self.sandbox.run('breakpoint()') + assert not result['success'] + + def test_reject_compile(self): + result = self.sandbox.run('compile("x=1", "", "exec")') + assert not result['success'] + + def test_reject_globals(self): + result = self.sandbox.run('result = globals()') + assert not result['success'] + + def test_reject_importlib(self): + result = self.sandbox.run('import importlib') + assert not result['success'] + + def test_reject_shutil(self): + result = self.sandbox.run('import shutil') + assert not result['success'] + + def test_reject_ctypes(self): + result = self.sandbox.run('import ctypes') + assert not result['success'] + + def test_reject_multiprocessing(self): + result = self.sandbox.run('import multiprocessing') + assert not result['success'] + + def test_reject_threading(self): + result = self.sandbox.run('import threading') + assert not result['success'] + + def test_reject_http(self): + result = self.sandbox.run('import http') + assert not result['success'] + + def test_reject_urllib(self): + result = self.sandbox.run('import urllib') + assert not result['success'] + + +# ═══════════════════════════════════════════════ +# 4. 超时测试 +# ═══════════════════════════════════════════════ + +class TestSandboxTimeout: + """超时机制测试""" + + def test_timeout_infinite_loop(self): + sandbox = Sandbox(timeout=2, mode='threaded') + result = sandbox.run('while True: pass') + assert not result['success'] + assert result['killed'] + sandbox.cleanup() + + def test_timeout_sleep(self): + sandbox = Sandbox(timeout=2, mode='threaded') + result = sandbox.run('import time\ntime.sleep(30)\nresult = "done"') + # Should either timeout or fail on import + assert not result['success'] or result['killed'] + sandbox.cleanup() + + +# ═══════════════════════════════════════════════ +# 5. SandboxResult 测试 +# ═══════════════════════════════════════════════ + +class TestSandboxResult: + """SandboxResult 数据类测试""" + + def test_success_result(self): + r = SandboxResult(success=True, result="42", time=0.1) + assert r.success + d = r.to_dict() + assert d['success'] + assert d['result'] == '42' + assert d['time'] == 0.1 + + def test_failure_result(self): + r = SandboxResult(success=False, result="error", security_violation="blocked os") + assert not r.success + d = r.to_dict() + assert not d['success'] + assert d['security_violation'] == 'blocked os' + + def test_str_representation(self): + r = SandboxResult(success=True, result="ok", time=0.05) + s = str(r) + assert 'OK' in s + assert '0.050' in s + + +# ═══════════════════════════════════════════════ +# 6. 沙箱模式对比测试 +# ═══════════════════════════════════════════════ + +class TestSandboxModes: + """threaded vs subprocess 模式对比""" + + def test_threaded_mode(self): + sandbox = Sandbox(timeout=10, mode='threaded') + result = sandbox.run('result = 2 + 3') + assert result['success'] + assert '5' in result['result'] + sandbox.cleanup() + + def test_subprocess_mode(self): + sandbox = Sandbox(timeout=10, mode='subprocess') + result = sandbox.run('result = 2 + 3') + assert result['success'] + assert '5' in result['result'] + sandbox.cleanup() + + def test_both_modes_same_result(self): + code = 'import json\nresult = json.dumps({"x": 42})' + + t = Sandbox(timeout=10, mode='threaded') + s = Sandbox(timeout=10, mode='subprocess') + + r_threaded = t.run(code) + r_subprocess = s.run(code) + + assert r_threaded['success'] == r_subprocess['success'] + assert r_threaded['result'] == r_subprocess['result'] + + t.cleanup() + s.cleanup() + + +# ═══════════════════════════════════════════════ +# 7. 禁用沙箱测试 +# ═══════════════════════════════════════════════ + +class TestSandboxDisabled: + """沙箱禁用时的行为""" + + def test_disabled_rejects_all(self): + sandbox = Sandbox(timeout=10, mode='threaded') + sandbox.enabled = False + result = sandbox.run('result = 1') + assert not result['success'] + assert 'disabled' in result['result'].lower() or 'off' in result['security_violation'].lower() + sandbox.cleanup() + + +# ═══════════════════════════════════════════════ +# 8. 历史记录测试 +# ═══════════════════════════════════════════════ + +class TestSandboxHistory: + """沙箱执行历史记录""" + + def test_history_recorded(self): + sandbox = Sandbox(timeout=10, mode='threaded') + sandbox.run('result = 1') + sandbox.run('result = 2') + assert len(sandbox.history) == 2 + assert sandbox.history[0]['success'] + assert sandbox.history[1]['success'] + sandbox.cleanup() + + def test_failed_history_recorded(self): + sandbox = Sandbox(timeout=10, mode="threaded") + result = sandbox.run("x = 1 / 0") # runtime error, not security + assert not result["success"] + assert len(sandbox.history) == 1 + assert not sandbox.history[0]["success"] + sandbox.cleanup() + + def test_security_violation_no_history(self): + sandbox = Sandbox(timeout=10, mode="threaded") + result = sandbox.run("import os") # security violation + assert not result["success"] + # security violations may not record history + sandbox.cleanup() + + +# ═══════════════════════════════════════════════ +# 9. 变量注入安全测试 +# ═══════════════════════════════════════════════ + +class TestSandboxVariables: + """变量注入测试""" + + def test_inject_numbers(self): + sandbox = Sandbox(timeout=10, mode='threaded') + result = sandbox.run('result = a + b', variables={'a': 100, 'b': 200}) + assert result['success'] + assert '300' in result['result'] + sandbox.cleanup() + + def test_inject_strings(self): + sandbox = Sandbox(timeout=10, mode='threaded') + result = sandbox.run('result = name.upper()', variables={'name': 'hello'}) + assert result['success'] + assert 'HELLO' in result['result'] + sandbox.cleanup() + + def test_inject_dict(self): + sandbox = Sandbox(timeout=10, mode='threaded') + result = sandbox.run( + 'result = data["key"]', + variables={'data': {'key': 'found'}} + ) + assert result['success'] + assert 'found' in result['result'] + sandbox.cleanup() + + def test_inject_list(self): + sandbox = Sandbox(timeout=10, mode='threaded') + result = sandbox.run( + 'result = sum(items)', + variables={'items': [1, 2, 3, 4, 5]} + ) + assert result['success'] + assert '15' in result['result'] + sandbox.cleanup() diff --git a/xiaoli/sandbox.py b/xiaoli/sandbox.py index 256ceb3..8c75e34 100644 --- a/xiaoli/sandbox.py +++ b/xiaoli/sandbox.py @@ -1,100 +1,475 @@ -"""安全代码沙箱""" -import io +""" +安全代码沙箱 v2 - 多层隔离执行环境 + +安全层级: + 1. AST 静态分析 - 拒绝危险语法 + 2. 子进程隔离 - 独立进程执行,资源限制 + 3. 文件系统限制 - 仅允许临时目录读写 + 4. 网络阻断 - 禁止 socket 连接 + 5. 模块白名单 - 仅允许安全模块导入 + 6. 输出截断 - 防止内存耗尽 +""" import ast +import io +import json +import os import re +import shutil +import subprocess +import sys +import tempfile +import textwrap import time -import threading import traceback +import threading from contextlib import redirect_stdout, redirect_stderr -from colorama import Fore, Style +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Dict, List, Optional, Set, Tuple + import logging logger = logging.getLogger(__name__) -DANGEROUS_MODULES = frozenset({ +try: + import resource + HAS_RESOURCE = True +except ImportError: + HAS_RESOURCE = False + + +# ────────────────────────────────────────────── +# 安全策略配置 +# ────────────────────────────────────────────── + +BLOCKED_MODULES: Set[str] = frozenset({ 'os', 'sys', 'subprocess', 'shutil', 'socket', 'pickle', 'marshal', 'ctypes', 'multiprocessing', 'threading', 'signal', 'resource', 'posix', 'nt', 'builtins', - 'importlib', 'code', 'codeop', + 'importlib', 'code', 'codeop', 'bdb', 'pdb', + 'compileall', 'py_compile', 'zipimport', 'pkgutil', + 'http', 'urllib', 'ftplib', 'smtplib', 'xmlrpc', + 'asyncio', 'concurrent', '_thread', }) -DANGEROUS_NAMES = frozenset({ +BLOCKED_NAMES: Set[str] = frozenset({ 'exec', 'eval', 'compile', 'execfile', 'open', 'globals', 'locals', 'vars', 'getattr', 'setattr', 'delattr', '__import__', '__builtins__', '__subclasses__', + '__loader__', '__spec__', '__file__', '__name__', + 'breakpoint', 'exit', 'quit', 'help', }) -DANGEROUS_PATTERNS = [ +BLOCKED_AST_PATTERNS: List[str] = [ r'__', r'\beval\b', r'\bexec\b', r'\bcompile\b', r'\bglobals\b', r'\blocals\b', r'\bgetattr\b', - r'\bsetattr\b', r'\b__import__\b', + r'\bsetattr\b', r'\b__import__\b', r'\bbreakpoint\b', ] +DEFAULT_WHITELIST: List[str] = [ + 'json', 'math', 'datetime', 'random', 'collections', + 'itertools', 'functools', 'operator', 'typing', 're', + 'string', 'textwrap', 'copy', 'decimal', 'fractions', + 'statistics', 'bisect', 'heapq', 'array', 'struct', + 'hashlib', 'hmac', 'secrets', 'base64', 'binascii', + 'uuid', 'pathlib', 'enum', 'dataclasses', 'contextlib', +] -class Sandbox: - """安全代码执行沙箱""" +DEFAULT_MEMORY_LIMIT_MB = 256 +DEFAULT_CPU_TIME_LIMIT = 30 +DEFAULT_OUTPUT_LIMIT = 1_000_000 - def __init__(self, timeout=30, whitelist=None): - self.enabled = True - self.timeout = timeout - self.whitelist = whitelist or [ - 'json', 'math', 'datetime', 'random', 'collections', - 'itertools', 'functools', 'operator', 'typing', 're' - ] - self.history = [] - def is_safe(self, code: str) -> bool: - """快速安全检查""" - for pattern in DANGEROUS_PATTERNS: +# ────────────────────────────────────────────── +# 结果数据类 +# ────────────────────────────────────────────── + +@dataclass +class SandboxResult: + """沙箱执行结果""" + success: bool + result: str = "" + stdout: str = "" + stderr: str = "" + time: float = 0.0 + exit_code: int = 0 + killed: bool = False + security_violation: str = "" + + def to_dict(self) -> dict: + return { + 'success': self.success, + 'result': self.result, + 'stdout': self.stdout, + 'stderr': self.stderr, + 'time': round(self.time, 4), + 'exit_code': self.exit_code, + 'killed': self.killed, + 'security_violation': self.security_violation, + } + + def __str__(self) -> str: + status = "OK" if self.success else "FAIL" + parts = [f"{status}: {self.result}"] + if self.stdout: + parts.append(f"stdout={self.stdout[:200]}") + if self.stderr: + parts.append(f"stderr={self.stderr[:200]}") + if self.security_violation: + parts.append(f"security={self.security_violation}") + parts.append(f"{self.time:.3f}s") + return " | ".join(parts) + + +# ────────────────────────────────────────────── +# AST 安全检查器 +# ────────────────────────────────────────────── + +class ASTSecurityChecker(ast.NodeVisitor): + """AST 级静态安全分析""" + + def __init__(self, blocked_modules: Set[str] = None, + blocked_names: Set[str] = None): + self.blocked = blocked_modules or BLOCKED_MODULES + self.blocked_names = blocked_names or BLOCKED_NAMES + self.violations: List[str] = [] + + def _check_name(self, name: str, node: ast.AST): + lineno = getattr(node, 'lineno', 0) + if name in self.blocked_names: + self.violations.append(f"blocked name: {name} (line {lineno})") + if name.split('.')[0] in self.blocked: + self.violations.append(f"blocked module: {name} (line {lineno})") + + def visit_Import(self, node): + for alias in node.names: + self._check_name(alias.name, node) + self.generic_visit(node) + + def visit_ImportFrom(self, node): + if node.module: + self._check_name(node.module, node) + self.generic_visit(node) + + def visit_Call(self, node): + if isinstance(node.func, ast.Name): + self._check_name(node.func.id, node) + elif isinstance(node.func, ast.Attribute): + if isinstance(node.func.value, ast.Name): + name = node.func.value.id + if name in self.blocked_names: + lineno = getattr(node, 'lineno', 0) + self.violations.append( + f"blocked call: {name}.{node.func.attr} (line {lineno})" + ) + self.generic_visit(node) + + def visit_Name(self, node): + self._check_name(node.id, node) + self.generic_visit(node) + + def visit_Attribute(self, node): + if node.attr.startswith('__') and node.attr.endswith('__'): + if node.attr not in ('__init__', '__str__', '__repr__', + '__len__', '__contains__', '__iter__', + '__eq__', '__ne__', '__lt__', '__gt__', + '__le__', '__ge__', '__hash__', '__bool__', + '__add__', '__sub__', '__mul__', '__truediv__', + '__floordiv__', '__mod__', '__pow__', + '__enter__', '__exit__'): + lineno = getattr(node, 'lineno', 0) + self.violations.append( + f"blocked dunder: {node.attr} (line {lineno})" + ) + self.generic_visit(node) + + def visit_Delete(self, node): + self.violations.append(f"blocked del (line {getattr(node, 'lineno', 0)})") + self.generic_visit(node) + + def visit_Global(self, node): + self.violations.append(f"blocked global (line {getattr(node, 'lineno', 0)})") + self.generic_visit(node) + + def visit_Nonlocal(self, node): + self.violations.append(f"blocked nonlocal (line {getattr(node, 'lineno', 0)})") + self.generic_visit(node) + + def check(self, code: str) -> Tuple[bool, List[str]]: + """检查代码安全性,返回 (safe, violations)""" + self.violations = [] + + for pattern in BLOCKED_AST_PATTERNS: if re.search(pattern, code, re.IGNORECASE): - return False + self.violations.append(f"regex match: {pattern}") try: tree = ast.parse(code) - except SyntaxError: - return False + except SyntaxError as e: + return False, [f"syntax error: {e}"] - class Checker(ast.NodeVisitor): - def __init__(self): - self.ok = True + self.visit(tree) + return len(self.violations) == 0, self.violations - def _check(self, name): - if name in DANGEROUS_NAMES or name in DANGEROUS_MODULES: - self.ok = False - def visit_Import(self, node): - for alias in node.names: - if alias.name.split('.')[0] in DANGEROUS_MODULES: - self.ok = False - self.generic_visit(node) +# ────────────────────────────────────────────── +# 子进程包装脚本 +# ────────────────────────────────────────────── - def visit_ImportFrom(self, node): - if node.module and node.module.split('.')[0] in DANGEROUS_MODULES: - self.ok = False - self.generic_visit(node) +def _make_worker_script(code: str, variables_json: str, + whitelist_json: str, work_dir: str) -> str: + return textwrap.dedent(f'''\ +import sys +import os +import json +import io +import traceback +from contextlib import redirect_stdout, redirect_stderr - def visit_Call(self, node): - if isinstance(node.func, ast.Name): - self._check(node.func.id) - self.generic_visit(node) +os.chdir({work_dir!r}) - def visit_Name(self, node): - self._check(node.id) - self.generic_visit(node) +# Block networking +class _FakeSocket: + def __getattr__(self, name): + raise OSError("network access blocked by sandbox") +class _FakeSocketModule: + AF_INET = 0 + SOCK_STREAM = 0 + def __getattr__(self, name): + raise OSError("network access blocked by sandbox") + def socket(self, *a, **kw): + raise OSError("network access blocked by sandbox") - checker = Checker() - checker.visit(tree) - return checker.ok +import types +_fake_mod = _FakeSocketModule() +sys.modules['socket'] = _fake_mod +sys.modules['_socket'] = _fake_mod + +WHITELIST = json.loads({whitelist_json!r}) +safe_modules = {{}} +for mod_name in WHITELIST: + try: + safe_modules[mod_name] = __import__(mod_name) + except ImportError: + pass + +safe_builtins = {{ + 'print': print, 'len': len, 'str': str, 'int': int, + 'float': float, 'list': list, 'dict': dict, 'tuple': tuple, + 'set': set, 'frozenset': frozenset, + 'range': range, 'enumerate': enumerate, 'zip': zip, + 'sorted': sorted, 'reversed': reversed, + 'sum': sum, 'max': max, 'min': min, 'abs': abs, + 'round': round, 'pow': pow, 'divmod': divmod, + 'type': type, 'isinstance': isinstance, 'issubclass': issubclass, + 'hasattr': hasattr, 'callable': callable, 'id': id, + 'chr': chr, 'ord': ord, 'hex': hex, 'oct': oct, 'bin': bin, + 'bool': bool, 'bytes': bytes, 'bytearray': bytearray, + 'map': map, 'filter': filter, + 'any': any, 'all': all, + 'True': True, 'False': False, 'None': None, +}} + +user_vars = json.loads({variables_json!r}) + +exec_vars = {{}} +exec_vars.update(safe_builtins) +exec_vars.update(safe_modules) +exec_vars.update(user_vars) + +stdout_buf = io.StringIO() +stderr_buf = io.StringIO() +result_data = {{"success": True, "result": "", "stdout": "", "stderr": "", "error": ""}} + +try: + with redirect_stdout(stdout_buf), redirect_stderr(stderr_buf): + exec({code!r}, exec_vars) + val = exec_vars.get('result', None) + result_data["result"] = str(val) if val is not None else "executed ok" +except Exception as e: + result_data["success"] = False + result_data["error"] = str(e) + result_data["stderr"] = traceback.format_exc() + +result_data["stdout"] = stdout_buf.getvalue() +result_data["stderr"] = stderr_buf.getvalue() or result_data.get("stderr", "") + +max_out = 500000 +if len(result_data["stdout"]) > max_out: + result_data["stdout"] = result_data["stdout"][:max_out] + "\\n... (truncated)" +if len(result_data["stderr"]) > max_out: + result_data["stderr"] = result_data["stderr"][:max_out] + "\\n... (truncated)" + +sys.stdout.write(json.dumps(result_data, ensure_ascii=False)) +''') + + +# ────────────────────────────────────────────── +# 主沙箱类 +# ────────────────────────────────────────────── + +class Sandbox: + """安全代码执行沙箱 v2 + + mode: + - 'subprocess': 子进程隔离(默认,最安全) + - 'threaded': 线程内 exec(兼容模式) + """ + + def __init__(self, timeout: int = 30, + memory_limit_mb: int = DEFAULT_MEMORY_LIMIT_MB, + whitelist: List[str] = None, + mode: str = 'subprocess', + work_dir: str = None): + self.enabled = True + self.timeout = timeout + self.memory_limit_mb = memory_limit_mb + self.whitelist = whitelist or DEFAULT_WHITELIST + self.mode = mode + self.work_dir = work_dir or tempfile.mkdtemp(prefix='xiaoli_sandbox_') + self.history: List[Dict] = [] + self._checker = ASTSecurityChecker() + os.makedirs(self.work_dir, exist_ok=True) + + def is_safe(self, code: str) -> Tuple[bool, List[str]]: + """静态安全检查(AST + 正则)""" + return self._checker.check(code) def run(self, code: str, variables: dict = None) -> dict: - """执行代码""" + """执行代码,返回结果字典(兼容旧接口)""" + result = self.execute(code, variables) + return result.to_dict() + + def execute(self, code: str, variables: dict = None) -> SandboxResult: + """执行代码,返回 SandboxResult""" if not self.enabled: - return {'success': False, 'result': '沙箱已禁用', 'stdout': '', 'stderr': '', 'time': 0} + return SandboxResult( + success=False, result="sandbox disabled", + security_violation="sandbox is off" + ) - if not self.is_safe(code): - return {'success': False, 'result': '代码包含危险操作', 'stdout': '', 'stderr': '', 'time': 0} + safe, violations = self.is_safe(code) + if not safe: + return SandboxResult( + success=False, + result="code failed security check", + security_violation="; ".join(violations[:5]) + ) - exec_vars = dict(variables or {}) + if self.mode == 'subprocess': + result = self._run_subprocess(code, variables) + else: + result = self._run_threaded(code, variables) + + self.history.append({ + 'code': code[:500], + 'success': result.success, + 'time': result.time, + 'timestamp': time.time(), + }) + return result + + def cleanup(self): + """清理临时目录""" + try: + if os.path.exists(self.work_dir) and 'xiaoli_sandbox_' in self.work_dir: + shutil.rmtree(self.work_dir, ignore_errors=True) + except Exception: + pass + + # ── subprocess ── + + def _run_subprocess(self, code: str, variables: dict = None) -> SandboxResult: + vars_json = json.dumps(variables or {}, ensure_ascii=False, default=str) + whitelist_json = json.dumps(self.whitelist) + script = _make_worker_script(code, vars_json, whitelist_json, self.work_dir) + start = time.time() + + try: + proc = subprocess.run( + [sys.executable, '-c', script], + capture_output=True, + timeout=self.timeout, + cwd=self.work_dir, + env=self._safe_env(), + preexec_fn=self._set_limits if HAS_RESOURCE else None, + ) + elapsed = time.time() - start + + stdout = proc.stdout.decode('utf-8', errors='replace')[:DEFAULT_OUTPUT_LIMIT] + stderr = proc.stderr.decode('utf-8', errors='replace')[:DEFAULT_OUTPUT_LIMIT] + + try: + data = json.loads(stdout) + return SandboxResult( + success=data.get('success', False), + result=data.get('result', ''), + stdout=data.get('stdout', ''), + stderr=data.get('stderr', ''), + time=elapsed, + exit_code=proc.returncode, + ) + except json.JSONDecodeError: + return SandboxResult( + success=proc.returncode == 0, + result=stdout[:2000] if stdout else "no output", + stdout=stdout, + stderr=stderr, + time=elapsed, + exit_code=proc.returncode, + ) + + except subprocess.TimeoutExpired: + return SandboxResult( + success=False, + result=f"timeout ({self.timeout}s)", + time=self.timeout, + killed=True, + ) + except Exception as e: + return SandboxResult( + success=False, + result=f"subprocess error: {e}", + time=time.time() - start, + ) + + def _safe_env(self) -> dict: + env = { + 'PATH': '/usr/bin:/bin', + 'LANG': 'en_US.UTF-8', + 'LC_ALL': 'en_US.UTF-8', + 'HOME': self.work_dir, + 'TMPDIR': self.work_dir, + 'TEMP': self.work_dir, + 'TMP': self.work_dir, + 'PYTHONNOUSERSITE': '1', + 'PYTHONDONTWRITEBYTECODE': '1', + } + for key in ('TERM', 'COLUMNS', 'LINES', 'SHELL'): + val = os.environ.get(key) + if val: + env[key] = val + return env + + @staticmethod + def _set_limits(): + if not HAS_RESOURCE: + return + try: + mem_bytes = DEFAULT_MEMORY_LIMIT_MB * 1024 * 1024 + resource.setrlimit(resource.RLIMIT_AS, (mem_bytes, mem_bytes)) + resource.setrlimit(resource.RLIMIT_CPU, ( + DEFAULT_CPU_TIME_LIMIT, DEFAULT_CPU_TIME_LIMIT + 5 + )) + resource.setrlimit(resource.RLIMIT_NPROC, (0, 0)) + resource.setrlimit(resource.RLIMIT_FSIZE, (50 * 1024 * 1024, 50 * 1024 * 1024)) + except (ValueError, OSError): + pass + + # ── threaded (compat) ── + + def _run_threaded(self, code: str, variables: dict = None) -> SandboxResult: + exec_vars = {} exec_vars.update({ 'print': print, 'len': len, 'str': str, 'int': int, 'float': float, 'list': list, 'dict': dict, 'tuple': tuple, @@ -103,43 +478,63 @@ class Sandbox: 'min': min, 'abs': abs, 'round': round, 'type': type, 'isinstance': isinstance, }) - - for mod in self.whitelist: + for mod_name in self.whitelist: try: - exec_vars[mod] = __import__(mod) + exec_vars[mod_name] = __import__(mod_name) except ImportError: pass + if variables: + exec_vars.update(variables) - stdout = io.StringIO() - stderr = io.StringIO() - start = time.time() + stdout_buf = io.StringIO() + stderr_buf = io.StringIO() result = {'success': True, 'error': None, 'tb': None} def _exec(): try: - exec(code, exec_vars) + with redirect_stdout(stdout_buf), redirect_stderr(stderr_buf): + exec(code, exec_vars) except Exception as e: result['success'] = False result['error'] = str(e) result['tb'] = traceback.format_exc() + start = time.time() thread = threading.Thread(target=_exec, daemon=True) thread.start() thread.join(timeout=self.timeout) elapsed = time.time() - start if thread.is_alive(): - return {'success': False, 'result': f'超时 ({self.timeout}s)', - 'stdout': stdout.getvalue(), 'stderr': stderr.getvalue(), 'time': self.timeout} + return SandboxResult( + success=False, + result=f"timeout ({self.timeout}s)", + stdout=stdout_buf.getvalue(), + stderr=stderr_buf.getvalue(), + time=self.timeout, + killed=True, + ) + + stdout = stdout_buf.getvalue() + stderr = stderr_buf.getvalue() if result['success']: val = exec_vars.get('result', None) - return {'success': True, - 'result': str(val) if val is not None else '执行成功(无返回值)', - 'stdout': stdout.getvalue(), 'stderr': stderr.getvalue(), 'time': elapsed} + return SandboxResult( + success=True, + result=str(val) if val is not None else "executed ok", + stdout=stdout, + stderr=stderr, + time=elapsed, + ) else: - return {'success': False, - 'result': f"错误: {result['error']}", - 'stdout': stdout.getvalue(), - 'stderr': stderr.getvalue() + '\n' + (result['tb'] or ''), - 'time': elapsed} + return SandboxResult( + success=False, + result=f"error: {result['error']}", + stdout=stdout, + stderr=stderr + ('\n' + (result['tb'] or '')), + time=elapsed, + ) + + def __del__(self): + self.cleanup()