From e50d5a1747b20a462ede1c55cedb07c8500dcfce Mon Sep 17 00:00:00 2001 From: Aaron Date: Sun, 18 Feb 2024 16:41:36 +0000 Subject: [PATCH] Command line code sanitation (#1627) * UPDATE - add commandline sanitation class, update local_commandline_code_executor.py and renamed test for code_utils.py * FIX - precommit run * UPDATE - moved sanitation function to LocalCommandlineCodeExecutor, moved testo to test_commandline_code_executor.py * UPDATE - added docstring notice to sanitize_code * Update autogen/coding/local_commandline_code_executor.py Co-authored-by: Eric Zhu * FIX - regular expression * FIX - function invocation in tests * UPDATE - pre-commit run * FIX - pre-commit run -_- --------- Co-authored-by: Ward Co-authored-by: Eric Zhu --- .../coding/local_commandline_code_executor.py | 28 +++++++++++++++++++ test/coding/test_commandline_code_executor.py | 19 +++++++++++++ test/{test_code.py => test_code_utils.py} | 0 3 files changed, 47 insertions(+) rename test/{test_code.py => test_code_utils.py} (100%) diff --git a/autogen/coding/local_commandline_code_executor.py b/autogen/coding/local_commandline_code_executor.py index 5989e940d7..fe3dbf8d31 100644 --- a/autogen/coding/local_commandline_code_executor.py +++ b/autogen/coding/local_commandline_code_executor.py @@ -1,4 +1,5 @@ import os +import re import uuid import warnings from typing import Any, ClassVar, List, Optional @@ -43,6 +44,8 @@ class LocalCommandlineCodeExecutor(BaseModel): the working directory, and a unique file is generated and saved in the working directory for each code block. The code blocks are executed in the order they are received. + Command line code is sanitized using regular expression match against a list of dangerous commands in order to prevent self-destructive + commands from being executed which may potentially affect the users environment. Currently the only supported languages is Python and shell scripts. For Python code, use the language "python" for the code block. For shell scripts, use the language "bash", "shell", or "sh" for the code @@ -108,6 +111,28 @@ If you want the user to save the code in a file before executing it, put # filen """(Experimental) Export a code extractor that can be used by an agent.""" return MarkdownCodeExtractor() + @staticmethod + def sanitize_command(lang: str, code: str) -> None: + """ + Sanitize the code block to prevent dangerous commands. + This approach acknowledges that while Docker or similar + containerization/sandboxing technologies provide a robust layer of security, + not all users may have Docker installed or may choose not to use it. + Therefore, having a baseline level of protection helps mitigate risks for users who, + either out of choice or necessity, run code outside of a sandboxed environment. + """ + dangerous_patterns = [ + (r"\brm\s+-rf\b", "Use of 'rm -rf' command is not allowed."), + (r"\bmv\b.*?\s+/dev/null", "Moving files to /dev/null is not allowed."), + (r"\bdd\b", "Use of 'dd' command is not allowed."), + (r">\s*/dev/sd[a-z][1-9]?", "Overwriting disk blocks directly is not allowed."), + (r":\(\)\{\s*:\|\:&\s*\};:", "Fork bombs are not allowed."), + ] + if lang in ["bash", "shell", "sh"]: + for pattern, message in dangerous_patterns: + if re.search(pattern, code): + raise ValueError(f"Potentially dangerous command detected: {message}") + def execute_code_blocks(self, code_blocks: List[CodeBlock]) -> CommandlineCodeResult: """(Experimental) Execute the code blocks and return the result. @@ -119,6 +144,9 @@ If you want the user to save the code in a file before executing it, put # filen logs_all = "" for i, code_block in enumerate(code_blocks): lang, code = code_block.language, code_block.code + + LocalCommandlineCodeExecutor.sanitize_command(lang, code) + print( colored( f"\n>>>>>>>> EXECUTING CODE BLOCK {i} (inferred language is {lang})...", diff --git a/test/coding/test_commandline_code_executor.py b/test/coding/test_commandline_code_executor.py index ba0853d731..5a9931890f 100644 --- a/test/coding/test_commandline_code_executor.py +++ b/test/coding/test_commandline_code_executor.py @@ -177,3 +177,22 @@ def _test_conversable_agent_code_execution(executor: CodeExecutor) -> None: sender=ConversableAgent("user", llm_config=False, code_execution_config=False), ) assert "hello extract code" in reply # type: ignore[operator] + + +# Test cases for dangerous commands that should be caught by the sanitizer +@pytest.mark.parametrize( + "lang, code, expected_message", + [ + ("bash", "rm -rf /", "Use of 'rm -rf' command is not allowed."), + ("bash", "mv myFile /dev/null", "Moving files to /dev/null is not allowed."), + ("bash", "dd if=/dev/zero of=/dev/sda", "Use of 'dd' command is not allowed."), + ("bash", "echo Hello > /dev/sda", "Overwriting disk blocks directly is not allowed."), + ("bash", ":(){ :|:& };:", "Fork bombs are not allowed."), + ], +) +def test_dangerous_commands(lang, code, expected_message): + with pytest.raises(ValueError) as exc_info: + LocalCommandlineCodeExecutor.sanitize_command(lang, code) + assert expected_message in str( + exc_info.value + ), f"Expected message '{expected_message}' not found in '{str(exc_info.value)}'" diff --git a/test/test_code.py b/test/test_code_utils.py similarity index 100% rename from test/test_code.py rename to test/test_code_utils.py