From c82936b1b18d500265d285540949640812c70f36 Mon Sep 17 00:00:00 2001 From: christop <825583681@qq.com> Date: Thu, 18 Jun 2026 06:09:02 +0800 Subject: [PATCH] Harden task/step output handlers against code execution in model output Three handlers passed the raw LLM reply through a cosmetic JsonCleaner.extract_json_array_section() slice and then straight to eval(): - TaskOutputHandler.handle() (superagi/agent/output_handler.py) - ReplaceTaskOutputHandler.handle() (superagi/agent/output_handler.py) - QueueStepHandler._process_reply() (superagi/agent/queue_step_handler.py) The cleaner only slices the substring between the first '[' and last ']' without validating its contents, so a model reply such as [__import__('os').system(...)] or [__import__('subprocess').getoutput(...)] was evaluated as Python and executed arbitrary code on the backend host. All three sinks ship on default-seeded paths: the TaskOutputHandler/ReplaceTaskOutputHandler branches back the default 'Dynamic Task Queue-I' iteration workflow, and the TASK_QUEUE tool step that reaches QueueStepHandler is seeded into the built-in Sales Engagement and Recruitment workflows. This is a single underlying vulnerability (unsafe eval of model output) reachable through three distinct handler call-sites. The fix adds one shared guard and wires it into all three: - Add JsonCleaner.parse_array_section(), which extracts the array section (reusing extract_json_array_section) and parses it with ast.literal_eval. Legitimate replies (a literal list of task strings, in either single- or double-quoted form as emitted by OpenAI) parse exactly as before, but embedded function calls/imports now raise ValueError instead of executing. Uses the stdlib ast module only; public signatures are unchanged. - TaskOutputHandler.handle(), ReplaceTaskOutputHandler.handle() and QueueStepHandler._process_reply() now call parse_array_section() instead of extract_json_array_section()+eval(). Adds regression tests for each call-site asserting that os.system / subprocess.getoutput payloads are rejected (ValueError) and never invoked, and that a legitimate literal array is still parsed and enqueued. Signed-off-by: christop <825583681@qq.com> --- superagi/agent/output_handler.py | 6 +-- superagi/agent/queue_step_handler.py | 6 +-- superagi/helper/json_cleaner.py | 34 +++++++++++++ tests/unit_tests/agent/test_output_handler.py | 49 +++++++++++++++++++ .../agent/test_queue_step_handler.py | 26 ++++++++++ 5 files changed, 114 insertions(+), 7 deletions(-) diff --git a/superagi/agent/output_handler.py b/superagi/agent/output_handler.py index 1fdeb75316..9b813e3e60 100644 --- a/superagi/agent/output_handler.py +++ b/superagi/agent/output_handler.py @@ -145,8 +145,7 @@ def __init__(self, agent_execution_id: int, agent_config: dict): self.agent_config = agent_config def handle(self, session, assistant_reply): - assistant_reply = JsonCleaner.extract_json_array_section(assistant_reply) - tasks = eval(assistant_reply) + tasks = JsonCleaner.parse_array_section(assistant_reply) tasks = np.array(tasks).flatten().tolist() for task in reversed(tasks): self.task_queue.add_task(task) @@ -176,8 +175,7 @@ def __init__(self, agent_execution_id: int, agent_config: dict): self.agent_config = agent_config def handle(self, session, assistant_reply): - assistant_reply = JsonCleaner.extract_json_array_section(assistant_reply) - tasks = eval(assistant_reply) + tasks = JsonCleaner.parse_array_section(assistant_reply) self.task_queue.clear_tasks() for task in reversed(tasks): self.task_queue.add_task(task) diff --git a/superagi/agent/queue_step_handler.py b/superagi/agent/queue_step_handler.py index fcd9baf1f4..dc6f8721b4 100644 --- a/superagi/agent/queue_step_handler.py +++ b/superagi/agent/queue_step_handler.py @@ -74,9 +74,9 @@ def _consume_from_queue(self, task_queue: TaskQueue): task_queue.complete_task("PROCESSED") def _process_reply(self, task_queue: TaskQueue, assistant_reply: str): - assistant_reply = JsonCleaner.extract_json_array_section(assistant_reply) - print("Queue reply:", assistant_reply) - task_array = np.array(eval(assistant_reply)).flatten().tolist() + tasks = JsonCleaner.parse_array_section(assistant_reply) + print("Queue reply:", tasks) + task_array = np.array(tasks).flatten().tolist() for task in task_array: task_queue.add_task(str(task)) logger.info("RAMRAM: Added task to queue: ", task) diff --git a/superagi/helper/json_cleaner.py b/superagi/helper/json_cleaner.py index 1b2cf30a92..31fe1a11a1 100644 --- a/superagi/helper/json_cleaner.py +++ b/superagi/helper/json_cleaner.py @@ -1,3 +1,4 @@ +import ast import json import re from superagi.lib.logger import logger @@ -7,6 +8,39 @@ class JsonCleaner: + @classmethod + def parse_array_section(cls, input_str: str = ""): + """ + Safely parse the array section of an LLM reply into a Python list. + + This is the safe replacement for ``eval()`` on model output. Model + replies are only ever expected to contain a literal array of task + strings (e.g. ``["task1", "task2"]`` or, because OpenAI emits + ``str(dict)``-style payloads, single-quoted ``['task1', 'task2']``). + ``ast.literal_eval`` accepts both quoting styles but, unlike ``eval``, + evaluates *only* Python literals, so it can never execute function + calls, imports or other code embedded in the reply + (e.g. ``[__import__('os').system('id')]`` raises instead of running). + + Args: + input_str (str): The (already array-sliced) model reply. + + Returns: + list: The parsed list of tasks. + + Raises: + ValueError: If the input is not a literal Python/JSON list. + """ + array_section = cls.extract_json_array_section(input_str) + try: + parsed = ast.literal_eval(array_section) + except (ValueError, SyntaxError, TypeError, MemoryError, RecursionError): + logger.error("JsonCleaner.parse_array_section: unable to parse array section safely") + raise ValueError("Model reply did not contain a valid literal array") + if not isinstance(parsed, (list, tuple)): + raise ValueError("Model reply did not contain a valid literal array") + return list(parsed) + @classmethod def clean_boolean(cls, input_str: str = ""): """ diff --git a/tests/unit_tests/agent/test_output_handler.py b/tests/unit_tests/agent/test_output_handler.py index 45474f217b..35efc06197 100644 --- a/tests/unit_tests/agent/test_output_handler.py +++ b/tests/unit_tests/agent/test_output_handler.py @@ -125,6 +125,30 @@ def test_task_output_handle_method(extract_json_array_section_mock, get_tasks_mo assert response.status == "PENDING" +# Security regression: TaskOutputHandler must not execute code embedded in the +# model reply. A prompt-injected/adversarial model can return a string that is a +# valid Python *expression* (e.g. a call to os.system); the handler must reject +# it as a non-literal instead of evaluating it. +@patch.object(TaskQueue, 'add_task') +@patch.object(TaskQueue, 'get_tasks') +def test_task_output_handle_rejects_code_execution(get_tasks_mock, add_task_mock): + agent_execution_id = 1 + agent_config = {"agent_id": 2} + # Payload that, under eval(), would invoke os.system and run a command. + malicious_reply = "[__import__('os').system('echo pwned')]" + get_tasks_mock.return_value = [] + handler = TaskOutputHandler(agent_execution_id, agent_config) + session_mock = MagicMock() + + with patch('os.system') as os_system_mock: + with pytest.raises(ValueError): + handler.handle(session_mock, malicious_reply) + os_system_mock.assert_not_called() + + # No tasks should have been queued from the malicious payload. + add_task_mock.assert_not_called() + + # Test for ReplaceTaskOutputHandler @patch.object(TaskQueue, 'clear_tasks') @patch.object(TaskQueue, 'add_task') @@ -152,3 +176,28 @@ def test_handle_method(extract_json_array_section_mock, get_tasks_mock, add_task assert add_task_mock.call_count == len(tasks) get_tasks_mock.assert_called_once() assert response.status == "PENDING" + + +# Security regression: ReplaceTaskOutputHandler must not execute code embedded in +# the model reply. A prompt-injected/adversarial model can return a string that is +# a valid Python *expression* (e.g. a call to subprocess.getoutput); the handler +# must reject it as a non-literal instead of evaluating it. +@patch.object(TaskQueue, 'clear_tasks') +@patch.object(TaskQueue, 'add_task') +@patch.object(TaskQueue, 'get_tasks') +def test_replace_task_output_handle_rejects_code_execution(get_tasks_mock, add_task_mock, clear_tasks_mock): + agent_execution_id = 1 + agent_config = {} + # Payload that, under eval(), would invoke subprocess.getoutput and run a command. + malicious_reply = "[__import__('subprocess').getoutput('echo pwned')]" + get_tasks_mock.return_value = [] + handler = ReplaceTaskOutputHandler(agent_execution_id, agent_config) + session_mock = MagicMock() + + with patch('subprocess.getoutput') as getoutput_mock: + with pytest.raises(ValueError): + handler.handle(session_mock, malicious_reply) + getoutput_mock.assert_not_called() + + # No tasks should have been queued from the malicious payload. + add_task_mock.assert_not_called() diff --git a/tests/unit_tests/agent/test_queue_step_handler.py b/tests/unit_tests/agent/test_queue_step_handler.py index 01f7da6ee8..c6594964df 100644 --- a/tests/unit_tests/agent/test_queue_step_handler.py +++ b/tests/unit_tests/agent/test_queue_step_handler.py @@ -60,6 +60,32 @@ def test_add_to_queue(task_queue_mock, agent_execution_feed_mock, queue_step_han queue_step_handler._process_reply.assert_called_once_with(task_queue_mock, '{"reply": ["task1", "task2"]}') +# Security regression: _process_reply must not execute code embedded in the model +# reply. A prompt-injected/adversarial model can return a string that is a valid +# Python *expression* (e.g. a call to os.system); the queue step must reject it as a +# non-literal instead of evaluating it. +def test_process_reply_rejects_code_execution(queue_step_handler): + task_queue_mock = Mock() + # Payload that, under eval(), would invoke os.system and run a command. + malicious_reply = "[__import__('os').system('echo pwned')]" + + with patch('os.system') as os_system_mock: + with pytest.raises(ValueError): + queue_step_handler._process_reply(task_queue_mock, malicious_reply) + os_system_mock.assert_not_called() + + # No task should have been added from the malicious payload. + task_queue_mock.add_task.assert_not_called() + + +def test_process_reply_parses_literal_array(queue_step_handler): + # A legitimate model reply (a literal list of task strings) must still be parsed + # and enqueued, including the single-quoted str(dict)-style form OpenAI emits. + task_queue_mock = Mock() + queue_step_handler._process_reply(task_queue_mock, "['task1', 'task2']") + assert task_queue_mock.add_task.call_count == 2 + + @patch("superagi.agent.queue_step_handler.TaskQueue") @patch("superagi.agent.queue_step_handler.AgentExecutionFeed") def test_consume_from_queue(task_queue_mock, agent_execution_feed_mock, queue_step_handler, step_tool):