From 466ce2375c5638bbf1e498c71a8d9b47ffbb5b04 Mon Sep 17 00:00:00 2001
From: Max Novich <maksymstepanenko1990@gmail.com>
Date: Tue, 3 Sep 2024 21:10:08 -0700
Subject: [PATCH] added some regex based checks for dangerous commands (#38)

---
 src/goose/toolkit/developer.py          |  7 +++---
 src/goose/utils/check_shell_command.py  | 33 +++++++++++++++++++++++++
 tests/utils/test_check_shell_command.py | 27 ++++++++++++++++++++
 3 files changed, 64 insertions(+), 3 deletions(-)
 create mode 100644 src/goose/utils/check_shell_command.py
 create mode 100644 tests/utils/test_check_shell_command.py

diff --git a/src/goose/toolkit/developer.py b/src/goose/toolkit/developer.py
index 6f53ca4d..e977d13d 100644
--- a/src/goose/toolkit/developer.py
+++ b/src/goose/toolkit/developer.py
@@ -1,6 +1,7 @@
 from pathlib import Path
 from subprocess import CompletedProcess, run
 from typing import List
+from goose.utils.check_shell_command import is_dangerous_command
 
 from exchange import Message
 from rich import box
@@ -135,7 +136,7 @@ class Developer(Toolkit):
             command (str): The shell command to run. It can support multiline statements
                 if you need to run more than one at a time
         """
-        self.notifier.status("running shell command")
+        self.notifier.status("planning to run shell command")
         # Log the command being executed in a visually structured format (Markdown).
         # The `.log` method is used here to log the command execution in the application's UX
         # this method is dynamically attached to functions in the Goose framework to handle user-visible
@@ -156,13 +157,13 @@ class Developer(Toolkit):
             rating = int(rating)
         except ValueError:
             rating = 5  # if we can't interpret we default to unsafe
-        if int(rating) > 3:
+        if is_dangerous_command(command) or int(rating) > 3:
             if not keep_unsafe_command_prompt(command):
                 raise RuntimeError(
                     f"The command {command} was rejected as dangerous by the user."
                     + " Do not proceed further, instead ask for instructions."
                 )
-
+        self.notifier.status("running shell command")
         result: CompletedProcess = run(command, shell=True, text=True, capture_output=True, check=False)
         if result.returncode == 0:
             output = "Command succeeded"
diff --git a/src/goose/utils/check_shell_command.py b/src/goose/utils/check_shell_command.py
new file mode 100644
index 00000000..6081de30
--- /dev/null
+++ b/src/goose/utils/check_shell_command.py
@@ -0,0 +1,33 @@
+import re
+
+
+def is_dangerous_command(command: str) -> bool:
+    """
+    Check if the command matches any dangerous patterns.
+
+    Args:
+        command (str): The shell command to check.
+
+    Returns:
+        bool: True if the command is dangerous, False otherwise.
+    """
+    dangerous_patterns = [
+        r"\brm\b",  # rm command
+        r"\bgit\s+push\b",  # git push command
+        r"\bsudo\b",  # sudo command
+        # Add more dangerous command patterns here
+        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
+        # Manipulating files in ~/ directly or dot files
+        r"^~/[^/]+$",  # Files directly in home directory
+        r"/\.[^/]+$",  # Dot files
+    ]
+    for pattern in dangerous_patterns:
+        if re.search(pattern, command):
+            return True
+    return False
diff --git a/tests/utils/test_check_shell_command.py b/tests/utils/test_check_shell_command.py
new file mode 100644
index 00000000..7d94a46d
--- /dev/null
+++ b/tests/utils/test_check_shell_command.py
@@ -0,0 +1,27 @@
+import pytest
+from goose.utils.check_shell_command import is_dangerous_command
+
+
+@pytest.mark.parametrize(
+    "command",
+    [
+        "rm -rf /",
+        "git push origin master",
+        "sudo reboot",
+        "mv /etc/passwd /tmp/",
+        "chmod 777 /etc/passwd",
+        "chown root:root /etc/passwd",
+        "mkfs -t ext4 /dev/sda1",
+        "systemctl stop nginx",
+        "reboot",
+        "shutdown now",
+        "echo hello > ~/.bashrc",
+    ],
+)
+def test_dangerous_commands(command):
+    assert is_dangerous_command(command)
+
+
+@pytest.mark.parametrize("command", ["ls -la", 'echo "Hello World"', "cp ~/folder/file.txt /tmp/"])
+def test_safe_commands(command):
+    assert not is_dangerous_command(command)