diff --git a/docs/reference/index.md b/docs/reference/index.md index cbcb8e4d..865d3816 100644 --- a/docs/reference/index.md +++ b/docs/reference/index.md @@ -35,4 +35,4 @@ - [goose.utils.ask](goose/utils/ask.md) - [goose.utils.file_utils](goose/utils/file_utils.md) - [goose.utils.session_file](goose/utils/session_file.md) -- [goose.utils.shell.is_dangerous_command](goose/utils/shell/is_dangerous_command.md) +- [goose.utils.command_checker.is_dangerous_command](goose/utils/command_checker/is_dangerous_command.md) diff --git a/src/goose/synopsis/process_manager.py b/src/goose/synopsis/process_manager.py index cb112d89..7474e990 100644 --- a/src/goose/synopsis/process_manager.py +++ b/src/goose/synopsis/process_manager.py @@ -7,7 +7,8 @@ from goose.synopsis.system import system from goose.synopsis.util import log_command from goose.toolkit.utils import RULEPREFIX, RULESTYLE -from goose.utils.shell import is_dangerous_command, keep_unsafe_command_prompt +from goose.utils.command_checker import is_dangerous_command +from goose.utils.shell import keep_unsafe_command_prompt ProcessManagerCommand = Literal["start", "list", "view_output", "cancel"] diff --git a/src/goose/utils/command_checker.py b/src/goose/utils/command_checker.py new file mode 100644 index 00000000..b495c4b3 --- /dev/null +++ b/src/goose/utils/command_checker.py @@ -0,0 +1,49 @@ +import re +from typing import List + +_dangerous_patterns = [ + # Commands that are generally unsafe + r"\brm\b", # rm command + r"\bgit\s+push\b", # git push command + r"\bsudo\b", # sudo command + r"\bmv\b", # mv command + r"\bchmod\b", # chmod command + r"\bchown\b", # chown command + r"\bmkfs\b", # mkfs command + r"\bsystemctl\b", # systemctl command + r"\breboot\b", # reboot command + r"\bshutdown\b", # shutdown command + # Commands that kill processes + r"\b(kill|pkill|killall|xkill|skill)\b", + r"\bfuser\b\s*-[kK]", # fuser -k command + # Target files that are unsafe + r"\b~\/\.|\/\.\w+", # commands that point to files or dirs in home that start with a dot (dotfiles) +] +_compiled_patterns = [re.compile(pattern) for pattern in _dangerous_patterns] + + +def is_dangerous_command(command: str) -> bool: + """ + Check if the command matches any dangerous patterns. + + Dangerous patterns in this function are defined as commands that may present risk to system stability. + + Args: + command (str): The shell command to check. + + Returns: + bool: True if the command is dangerous, False otherwise. + """ + return any(pattern.search(command) for pattern in _compiled_patterns) + + +def add_dangerous_command_patterns(patterns: List[str]) -> None: + """ + Add additional dangerous patterns to the command checker. Intended to be + called in plugins that add additional high-specificity dangerous commands. + + Args: + patterns (List[str]): The regex patterns to add to the dangerous patterns list. + """ + _dangerous_patterns.extend(patterns) + _compiled_patterns.extend([re.compile(pattern) for pattern in patterns]) diff --git a/src/goose/utils/shell.py b/src/goose/utils/shell.py index c1a078e2..d346dd8b 100644 --- a/src/goose/utils/shell.py +++ b/src/goose/utils/shell.py @@ -6,46 +6,11 @@ from goose.notifier import Notifier from goose.utils.ask import ask_an_ai +from goose.utils.command_checker import is_dangerous_command from goose.view import ExchangeView from rich.prompt import Confirm -def is_dangerous_command(command: str) -> bool: - """ - Check if the command matches any dangerous patterns. - - Dangerous patterns in this function are defined as commands that may present risk to system stability. - - Args: - command (str): The shell command to check. - - Returns: - bool: True if the command is dangerous, False otherwise. - """ - dangerous_patterns = [ - # Commands that are generally unsafe - r"\brm\b", # rm command - r"\bgit\s+push\b", # git push command - r"\bsudo\b", # sudo command - r"\bmv\b", # mv command - r"\bchmod\b", # chmod command - r"\bchown\b", # chown command - r"\bmkfs\b", # mkfs command - r"\bsystemctl\b", # systemctl command - r"\breboot\b", # reboot command - r"\bshutdown\b", # shutdown command - # Commands that kill processes - r"\b(kill|pkill|killall|xkill|skill)\b", - r"\bfuser\b\s*-[kK]", # fuser -k command - # Target files that are unsafe - r"\b~\/\.|\/\.\w+", # commands that point to files or dirs in home that start with a dot (dotfiles) - ] - for pattern in dangerous_patterns: - if re.search(pattern, command): - return True - return False - - def keep_unsafe_command_prompt(command: str) -> bool: message = f"\nWe flagged the command - [bold red]{command}[/] - as potentially unsafe, do you want to proceed?" return Confirm.ask(message, default=True) diff --git a/tests/utils/test_check_shell_command.py b/tests/utils/test_check_shell_command.py index 7484b193..08749f73 100644 --- a/tests/utils/test_check_shell_command.py +++ b/tests/utils/test_check_shell_command.py @@ -1,5 +1,5 @@ import pytest -from goose.utils.shell import is_dangerous_command +from goose.utils.command_checker import add_dangerous_command_patterns, is_dangerous_command @pytest.mark.parametrize( @@ -40,3 +40,11 @@ def test_dangerous_commands(command): ) def test_safe_commands(command): assert not is_dangerous_command(command) + + +def test_add_dangerous_patterns(): + add_dangerous_command_patterns(["echo hello"]) + assert is_dangerous_command("echo hello") + + # and that the original commands are still flagged + assert is_dangerous_command("rm -rf /")