Skip to content

Commit

Permalink
feat: allow plugins to set custom patterns for dangerous commands (#281)
Browse files Browse the repository at this point in the history
  • Loading branch information
ajgray-stripe authored Dec 10, 2024
1 parent 64a1ee4 commit db11a22
Show file tree
Hide file tree
Showing 5 changed files with 62 additions and 39 deletions.
2 changes: 1 addition & 1 deletion docs/reference/index.md
Original file line number Diff line number Diff line change
Expand Up @@ -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)
3 changes: 2 additions & 1 deletion src/goose/synopsis/process_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]

Expand Down
49 changes: 49 additions & 0 deletions src/goose/utils/command_checker.py
Original file line number Diff line number Diff line change
@@ -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])
37 changes: 1 addition & 36 deletions src/goose/utils/shell.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
10 changes: 9 additions & 1 deletion tests/utils/test_check_shell_command.py
Original file line number Diff line number Diff line change
@@ -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(
Expand Down Expand Up @@ -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 /")

0 comments on commit db11a22

Please sign in to comment.