Merge pull request #46 from plasma-umass/pr46

Add --unsafe mode, add barebone safe mode to LLDB/GDB.

src/chatdbg/chatdbg_gdb.py

@@ -13,6 +13,7 @@
     _SkippedFramesEntry,
 )
 from chatdbg.util.config import chatdbg_config
+from chatdbg.native_util.safety import command_is_safe
 
 # The file produced by the panic handler if the Rust program is using the chatdbg crate.
 RUST_PANIC_LOG_FILENAME = "panic_log.txt"
@@ -262,7 +263,7 @@ def _prompt_stack(self):
         """
         return None
 
-    def llm_debug(self, command: str) -> str:
+    def llm_debug(self, command: str):
         """
         {
             "name": "debug",
@@ -279,4 +280,6 @@ def llm_debug(self, command: str) -> str:
             }
         }
         """
+        if not chatdbg_config.unsafe and not command_is_safe(command):
+            return command, f"Command `{command}` is not allowed."
         return command, self._run_one_command(command)

src/chatdbg/chatdbg_lldb.py

@@ -13,6 +13,7 @@
     _SkippedFramesEntry,
 )
 from chatdbg.util.config import chatdbg_config
+from chatdbg.native_util.safety import command_is_safe
 
 # The file produced by the panic handler if the Rust program is using the chatdbg crate.
 RUST_PANIC_LOG_FILENAME = "panic_log.txt"
@@ -290,7 +291,7 @@ def _prompt_stack(self):
         """
         return None
 
-    def llm_debug(self, command: str) -> str:
+    def llm_debug(self, command: str):
         """
         {
             "name": "debug",
@@ -307,4 +308,6 @@ def llm_debug(self, command: str) -> str:
             }
         }
         """
+        if not chatdbg_config.unsafe and not command_is_safe(command):
+            return command, f"Command `{command}` is not allowed."
         return command, self._run_one_command(command)

src/chatdbg/chatdbg_pdb.py

@@ -285,7 +285,10 @@ def _getval(self, arg):
         Sandbox for evaluating expressions from the LLM.
         """
         try:
-            return sandbox_eval(arg, self.curframe.f_globals, self.curframe_locals)
+            if chatdbg_config.unsafe:
+                return super._getval(arg)
+            else:   
+                return sandbox_eval(arg, self.curframe.f_globals, self.curframe_locals)
         except NameError as e:
             self.error(f"NameError: {e}")
             return None

src/chatdbg/native_util/safety.py

@@ -0,0 +1,32 @@
+import re
+
+
+# A very simple whitelist-based approach.
+# If ChatDBG wants to call other commands not listed here, they should be
+# evaluated and added if not possibly harmful.
+def command_is_safe(cmd: str) -> bool:
+    cmd = cmd.strip()
+    command_name = cmd.split()[0]
+
+    # Allowed unconditionally.
+    if command_name in [
+        "apropos",
+        "bt",
+        "down",
+        "frame",
+        "h",
+        "help",
+        "language",
+        "l",
+        "list",
+        "source",
+        "up",
+        "version",
+    ]:
+        return True
+
+    # Allowed conditionally.
+    if command_name in ["p", "print"]:
+        return re.fullmatch(r"[a-zA-Z0-9_ *]*", cmd) is not None
+
+    return False

src/chatdbg/util/config.py

@@ -23,7 +23,7 @@ def _chatdbg_get_env(
     if type(default_value) == int:
         return int(v)
     elif type(default_value) == bool:
-        return v.lower() == "true"
+        return v.lower() == "true" or v.lower() == "1"
     else:
         return v
 
@@ -85,18 +85,23 @@ class ChatDBGConfig(Configurable):
 
     format = Unicode(
         _chatdbg_get_env("format", "md"),
-        help="The output format (text or md or md:simple or jupyter).",
+        help="The output format (text or md or md:simple or jupyter)",
     ).tag(config=True)
 
     instructions = Unicode(
         _chatdbg_get_env("instructions", ""),
-        help="The file for the initial instructions to the LLM, or '' for the default (possibly-model specific) version.",
+        help="The file for the initial instructions to the LLM, or '' for the default (possibly-model specific) version",
     ).tag(config=True)
 
     module_whitelist = Unicode(
         _chatdbg_get_env("module_whitelist", ""), help="The module whitelist file"
     ).tag(config=True)
 
+    unsafe = Bool(
+        _chatdbg_get_env("unsafe", False),
+        help="Disable any protections against GPT running harmful code or commands",
+    ).tag(config=True)
+
     _user_configurable = [
         debug,
         log,
@@ -105,6 +110,7 @@ class ChatDBGConfig(Configurable):
         no_stream,
         format,
         module_whitelist,
+        unsafe,
     ]
 
     def _parser(self):