Skip to content

Commit

Permalink
Command line code sanitation (#1627)
Browse files Browse the repository at this point in the history
* 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 <ekzhu@users.noreply.github.com>

* FIX - regular expression

* FIX - function invocation in tests

* UPDATE - pre-commit run

* FIX - pre-commit run -_-

---------

Co-authored-by: Ward <award40@LAMU0CLP74YXVX6.uhc.com>
Co-authored-by: Eric Zhu <ekzhu@users.noreply.github.com>
  • Loading branch information
3 people committed Feb 18, 2024
1 parent 5aee911 commit e50d5a1
Show file tree
Hide file tree
Showing 3 changed files with 47 additions and 0 deletions.
28 changes: 28 additions & 0 deletions autogen/coding/local_commandline_code_executor.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import os
import re
import uuid
import warnings
from typing import Any, ClassVar, List, Optional
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -108,6 +111,28 @@ def code_extractor(self) -> CodeExtractor:
"""(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.
Expand All @@ -119,6 +144,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

LocalCommandlineCodeExecutor.sanitize_command(lang, code)

print(
colored(
f"\n>>>>>>>> EXECUTING CODE BLOCK {i} (inferred language is {lang})...",
Expand Down
19 changes: 19 additions & 0 deletions test/coding/test_commandline_code_executor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)}'"
File renamed without changes.

0 comments on commit e50d5a1

Please sign in to comment.