From 9f1ae578faee2e322042b1df44761c875d44e2ac Mon Sep 17 00:00:00 2001 From: Ward Date: Sun, 11 Feb 2024 16:59:28 +0000 Subject: [PATCH 1/9] UPDATE - add commandline sanitation class, update local_commandline_code_executor.py and renamed test for code_utils.py --- autogen/code_utils.py | 22 +++++++++++++++++++ .../coding/local_commandline_code_executor.py | 7 +++++- test/{test_code.py => test_code_utils.py} | 14 ++++++++++++ 3 files changed, 42 insertions(+), 1 deletion(-) rename test/{test_code.py => test_code_utils.py} (95%) diff --git a/autogen/code_utils.py b/autogen/code_utils.py index 96c3f179779f..a9c8209f5906 100644 --- a/autogen/code_utils.py +++ b/autogen/code_utils.py @@ -222,6 +222,28 @@ def _cmd(lang): return "powershell" raise NotImplementedError(f"{lang} not recognized in code execution") +class CommandSanitizer: + """ + A utility to sanitize code blocks before execution to prevent running dangerous commands, + with specific messages for each pattern. + """ + dangerous_patterns: List[Tuple[str, str]] = [ + (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."), + ] + + @staticmethod + def sanitize(lang: str, code: str) -> None: + """ + Checks if the given code contains dangerous patterns and raises an appropriate error. + """ + if lang in ["bash", "shell", "sh"]: + for pattern, message in CommandSanitizer.dangerous_patterns: + if re.search(pattern, code): + raise ValueError(f"Potentially dangerous command: {message}") def is_docker_running() -> bool: """Check if docker is running. diff --git a/autogen/coding/local_commandline_code_executor.py b/autogen/coding/local_commandline_code_executor.py index 5989e940d7bc..7c1997f6557c 100644 --- a/autogen/coding/local_commandline_code_executor.py +++ b/autogen/coding/local_commandline_code_executor.py @@ -6,7 +6,7 @@ from pydantic import BaseModel, Field, field_validator from ..agentchat.agent import LLMAgent -from ..code_utils import execute_code +from ..code_utils import execute_code, CommandSanitizer from .base import CodeBlock, CodeExtractor, CodeResult from .markdown_code_extractor import MarkdownCodeExtractor @@ -43,6 +43,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 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 @@ -119,6 +121,9 @@ def execute_code_blocks(self, code_blocks: List[CodeBlock]) -> CommandlineCodeRe logs_all = "" for i, code_block in enumerate(code_blocks): lang, code = code_block.language, code_block.code + + CommandSanitizer.sanitize(lang, code) + print( colored( f"\n>>>>>>>> EXECUTING CODE BLOCK {i} (inferred language is {lang})...", diff --git a/test/test_code.py b/test/test_code_utils.py similarity index 95% rename from test/test_code.py rename to test/test_code_utils.py index 6a5a662ae43b..e5f522176aa8 100644 --- a/test/test_code.py +++ b/test/test_code_utils.py @@ -17,6 +17,7 @@ in_docker_container, decide_use_docker, check_can_use_docker_or_throw, + CommandSanitizer ) from conftest import skip_docker @@ -552,6 +553,19 @@ def test_non_dict_in_list(self): with self.assertRaises(TypeError): content_str(content) +# 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: + CommandSanitizer.sanitize(lang, code) + assert expected_message in str(exc_info.value), f"Expected message '{expected_message}' not found in '{str(exc_info.value)}'" + if __name__ == "__main__": # test_infer_lang() From 40cc84529c71764cfe7844778b15be983897c7d9 Mon Sep 17 00:00:00 2001 From: Ward Date: Sun, 11 Feb 2024 17:24:26 +0000 Subject: [PATCH 2/9] FIX - precommit run --- autogen/code_utils.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/autogen/code_utils.py b/autogen/code_utils.py index a9c8209f5906..248feeecad97 100644 --- a/autogen/code_utils.py +++ b/autogen/code_utils.py @@ -222,17 +222,19 @@ def _cmd(lang): return "powershell" raise NotImplementedError(f"{lang} not recognized in code execution") + class CommandSanitizer: """ A utility to sanitize code blocks before execution to prevent running dangerous commands, with specific messages for each pattern. """ + dangerous_patterns: List[Tuple[str, str]] = [ - (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."), + (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."), ] @staticmethod @@ -245,6 +247,7 @@ def sanitize(lang: str, code: str) -> None: if re.search(pattern, code): raise ValueError(f"Potentially dangerous command: {message}") + def is_docker_running() -> bool: """Check if docker is running. From daad69dd55e310bd9b838deda18f638af4f2c49b Mon Sep 17 00:00:00 2001 From: Ward Date: Tue, 13 Feb 2024 23:21:45 +0000 Subject: [PATCH 3/9] UPDATE - moved sanitation function to LocalCommandlineCodeExecutor, moved testo to test_commandline_code_executor.py --- autogen/code_utils.py | 25 ------------------- .../coding/local_commandline_code_executor.py | 20 +++++++++++++-- test/coding/test_commandline_code_executor.py | 13 ++++++++++ test/test_code_utils.py | 15 ----------- 4 files changed, 31 insertions(+), 42 deletions(-) diff --git a/autogen/code_utils.py b/autogen/code_utils.py index 248feeecad97..96c3f179779f 100644 --- a/autogen/code_utils.py +++ b/autogen/code_utils.py @@ -223,31 +223,6 @@ def _cmd(lang): raise NotImplementedError(f"{lang} not recognized in code execution") -class CommandSanitizer: - """ - A utility to sanitize code blocks before execution to prevent running dangerous commands, - with specific messages for each pattern. - """ - - dangerous_patterns: List[Tuple[str, str]] = [ - (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."), - ] - - @staticmethod - def sanitize(lang: str, code: str) -> None: - """ - Checks if the given code contains dangerous patterns and raises an appropriate error. - """ - if lang in ["bash", "shell", "sh"]: - for pattern, message in CommandSanitizer.dangerous_patterns: - if re.search(pattern, code): - raise ValueError(f"Potentially dangerous command: {message}") - - def is_docker_running() -> bool: """Check if docker is running. diff --git a/autogen/coding/local_commandline_code_executor.py b/autogen/coding/local_commandline_code_executor.py index 7c1997f6557c..ef2e04930e46 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 @@ -6,7 +7,7 @@ from pydantic import BaseModel, Field, field_validator from ..agentchat.agent import LLMAgent -from ..code_utils import execute_code, CommandSanitizer +from ..code_utils import execute_code from .base import CodeBlock, CodeExtractor, CodeResult from .markdown_code_extractor import MarkdownCodeExtractor @@ -110,6 +111,21 @@ def code_extractor(self) -> CodeExtractor: """(Experimental) Export a code extractor that can be used by an agent.""" return MarkdownCodeExtractor() + @staticmethod + def sanitize_code(lang: str, code: str) -> None: + """Sanitize the code block to prevent dangerous commands.""" + dangerous_patterns = [ + (r'\brm\b.*\b-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]', "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. @@ -122,7 +138,7 @@ def execute_code_blocks(self, code_blocks: List[CodeBlock]) -> CommandlineCodeRe for i, code_block in enumerate(code_blocks): lang, code = code_block.language, code_block.code - CommandSanitizer.sanitize(lang, code) + LocalCommandlineCodeExecutor.sanitize_code(lang, code) print( colored( diff --git a/test/coding/test_commandline_code_executor.py b/test/coding/test_commandline_code_executor.py index 1ee5e6be337e..be4924e6fcf4 100644 --- a/test/coding/test_commandline_code_executor.py +++ b/test/coding/test_commandline_code_executor.py @@ -177,3 +177,16 @@ 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)}'" \ No newline at end of file diff --git a/test/test_code_utils.py b/test/test_code_utils.py index e5f522176aa8..7a2dc42c6d5a 100644 --- a/test/test_code_utils.py +++ b/test/test_code_utils.py @@ -17,7 +17,6 @@ in_docker_container, decide_use_docker, check_can_use_docker_or_throw, - CommandSanitizer ) from conftest import skip_docker @@ -553,20 +552,6 @@ def test_non_dict_in_list(self): with self.assertRaises(TypeError): content_str(content) -# 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: - CommandSanitizer.sanitize(lang, code) - assert expected_message in str(exc_info.value), f"Expected message '{expected_message}' not found in '{str(exc_info.value)}'" - - if __name__ == "__main__": # test_infer_lang() test_extract_code() From ceb5784b8c9816bba83b844798126f4da70e90bc Mon Sep 17 00:00:00 2001 From: Ward Date: Tue, 13 Feb 2024 23:30:18 +0000 Subject: [PATCH 4/9] UPDATE - added docstring notice to sanitize_code --- autogen/coding/local_commandline_code_executor.py | 10 +++++++++- test/test_code_utils.py | 1 + 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/autogen/coding/local_commandline_code_executor.py b/autogen/coding/local_commandline_code_executor.py index ef2e04930e46..b3ea012a9f9c 100644 --- a/autogen/coding/local_commandline_code_executor.py +++ b/autogen/coding/local_commandline_code_executor.py @@ -113,7 +113,15 @@ def code_extractor(self) -> CodeExtractor: @staticmethod def sanitize_code(lang: str, code: str) -> None: - """Sanitize the code block to prevent dangerous commands.""" + """ + 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\b.*\b-rf\b', "Use of 'rm -rf' command is not allowed."), (r'\bmv\b.*\s/dev/null', "Moving files to /dev/null is not allowed."), diff --git a/test/test_code_utils.py b/test/test_code_utils.py index 7a2dc42c6d5a..6a5a662ae43b 100644 --- a/test/test_code_utils.py +++ b/test/test_code_utils.py @@ -552,6 +552,7 @@ def test_non_dict_in_list(self): with self.assertRaises(TypeError): content_str(content) + if __name__ == "__main__": # test_infer_lang() test_extract_code() From 1c8fdb77284008f5e79d580aeea59511630a1326 Mon Sep 17 00:00:00 2001 From: Aaron Date: Tue, 13 Feb 2024 23:31:07 +0000 Subject: [PATCH 5/9] Update autogen/coding/local_commandline_code_executor.py Co-authored-by: Eric Zhu --- autogen/coding/local_commandline_code_executor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/autogen/coding/local_commandline_code_executor.py b/autogen/coding/local_commandline_code_executor.py index b3ea012a9f9c..53148e3d4b80 100644 --- a/autogen/coding/local_commandline_code_executor.py +++ b/autogen/coding/local_commandline_code_executor.py @@ -44,7 +44,7 @@ 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 in order to prevent self destructive + 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. From d54bde489f557a0e3f74cd6a5309466548bfcf7d Mon Sep 17 00:00:00 2001 From: Ward Date: Tue, 13 Feb 2024 23:48:36 +0000 Subject: [PATCH 6/9] FIX - regular expression --- .../coding/local_commandline_code_executor.py | 22 +++++++++---------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/autogen/coding/local_commandline_code_executor.py b/autogen/coding/local_commandline_code_executor.py index b3ea012a9f9c..3208b6be9bd8 100644 --- a/autogen/coding/local_commandline_code_executor.py +++ b/autogen/coding/local_commandline_code_executor.py @@ -115,19 +115,19 @@ def code_extractor(self) -> CodeExtractor: def sanitize_code(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. - + 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\b.*\b-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]', "Overwriting disk blocks directly is not allowed."), - (r':\(\)\{\s*:\|\:&\s*\};:', "Fork bombs are not allowed."), + (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: From 62fd4d29edc1d55ed224a048da50fdd27ade779a Mon Sep 17 00:00:00 2001 From: Ward Date: Sun, 18 Feb 2024 10:46:35 +0000 Subject: [PATCH 7/9] FIX - function invocation in tests --- autogen/coding/local_commandline_code_executor.py | 4 ++-- test/coding/test_commandline_code_executor.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/autogen/coding/local_commandline_code_executor.py b/autogen/coding/local_commandline_code_executor.py index 8a0f67280eb6..0594ea0c21a2 100644 --- a/autogen/coding/local_commandline_code_executor.py +++ b/autogen/coding/local_commandline_code_executor.py @@ -112,7 +112,7 @@ def code_extractor(self) -> CodeExtractor: return MarkdownCodeExtractor() @staticmethod - def sanitize_code(lang: str, code: str) -> None: + def sanitize_command(lang: str, code: str) -> None: """ Sanitize the code block to prevent dangerous commands. This approach acknowledges that while Docker or similar @@ -146,7 +146,7 @@ def execute_code_blocks(self, code_blocks: List[CodeBlock]) -> CommandlineCodeRe for i, code_block in enumerate(code_blocks): lang, code = code_block.language, code_block.code - LocalCommandlineCodeExecutor.sanitize_code(lang, code) + LocalCommandlineCodeExecutor.sanitize_command(lang, code) print( colored( diff --git a/test/coding/test_commandline_code_executor.py b/test/coding/test_commandline_code_executor.py index fba90d762c3a..6bcb5cd0a835 100644 --- a/test/coding/test_commandline_code_executor.py +++ b/test/coding/test_commandline_code_executor.py @@ -188,5 +188,5 @@ def _test_conversable_agent_code_execution(executor: CodeExecutor) -> None: ]) def test_dangerous_commands(lang, code, expected_message): with pytest.raises(ValueError) as exc_info: - LocalCommandlineCodeExecutor.sanitize_command(lang, code) + LocalCommandlineCodeExecutor.sanitize_code(lang, code) assert expected_message in str(exc_info.value), f"Expected message '{expected_message}' not found in '{str(exc_info.value)}'" \ No newline at end of file From 3e3da546c76c40b79802a47692c64f450c72081e Mon Sep 17 00:00:00 2001 From: Ward Date: Sun, 18 Feb 2024 10:52:05 +0000 Subject: [PATCH 8/9] UPDATE - pre-commit run --- autogen/coding/local_commandline_code_executor.py | 1 - test/coding/test_commandline_code_executor.py | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/autogen/coding/local_commandline_code_executor.py b/autogen/coding/local_commandline_code_executor.py index 0594ea0c21a2..fe3dbf8d3184 100644 --- a/autogen/coding/local_commandline_code_executor.py +++ b/autogen/coding/local_commandline_code_executor.py @@ -120,7 +120,6 @@ def sanitize_command(lang: str, code: str) -> None: 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."), diff --git a/test/coding/test_commandline_code_executor.py b/test/coding/test_commandline_code_executor.py index 6bcb5cd0a835..fba90d762c3a 100644 --- a/test/coding/test_commandline_code_executor.py +++ b/test/coding/test_commandline_code_executor.py @@ -188,5 +188,5 @@ def _test_conversable_agent_code_execution(executor: CodeExecutor) -> None: ]) def test_dangerous_commands(lang, code, expected_message): with pytest.raises(ValueError) as exc_info: - LocalCommandlineCodeExecutor.sanitize_code(lang, code) + 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)}'" \ No newline at end of file From 87c94489d02b68f8d4902e3cb9ab6f85139a1c15 Mon Sep 17 00:00:00 2001 From: Ward Date: Sun, 18 Feb 2024 11:00:00 +0000 Subject: [PATCH 9/9] FIX - pre-commit run -_- --- test/coding/test_commandline_code_executor.py | 22 ++++++++++++------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/test/coding/test_commandline_code_executor.py b/test/coding/test_commandline_code_executor.py index fba90d762c3a..5a9931890f10 100644 --- a/test/coding/test_commandline_code_executor.py +++ b/test/coding/test_commandline_code_executor.py @@ -178,15 +178,21 @@ def _test_conversable_agent_code_execution(executor: CodeExecutor) -> None: ) 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."), -]) +@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)}'" \ No newline at end of file + assert expected_message in str( + exc_info.value + ), f"Expected message '{expected_message}' not found in '{str(exc_info.value)}'"