Skip to content

Commit

Permalink
feat:whitelist some commands
Browse files Browse the repository at this point in the history
  • Loading branch information
0xbrayo committed Oct 20, 2024
1 parent ff2277c commit c3fe334
Showing 1 changed file with 17 additions and 5 deletions.
22 changes: 17 additions & 5 deletions gptme/tools/shell.py
Original file line number Diff line number Diff line change
Expand Up @@ -243,15 +243,27 @@ def execute_shell(
"""Executes a shell command and returns the output."""
shell = get_shell()
assert not args
whitelist_commands = ["ls", "stat", "cd", "cat", "pwd", "echo"]
whitelisted = True

cmd = code.strip()
if cmd.startswith("$ "):
cmd = cmd[len("$ ") :]

print_preview(cmd, "bash")
if not confirm("Run command?"):
yield Message("system", "Command not run")
return
#NOTE: This does not handle control flow words like if, for, while.
regex = r"(?:^|[|&;]|\|\||&&|\n)\s*([^\s|&;]+)"

for match in re.finditer(regex, cmd):
for group in match.groups():
if group and group not in whitelist_commands:
whitelisted = False
break

if not whitelisted:
print_preview(cmd, "bash")
if not confirm("Run command?"):
yield Message("system", "Command not run")
return

try:
returncode, stdout, stderr = shell.run(cmd)
Expand All @@ -261,7 +273,7 @@ def execute_shell(
stdout = _shorten_stdout(stdout.strip(), pre_tokens=2000, post_tokens=8000)
stderr = _shorten_stdout(stderr.strip(), pre_tokens=2000, post_tokens=2000)

msg = _format_block_smart("Ran command", cmd, lang="bash") + "\n\n"
msg = _format_block_smart(f"Ran {'whitelisted ' if whitelisted else ''}command", cmd, lang="bash") + "\n\n"
if stdout:
msg += _format_block_smart("", stdout, "stdout") + "\n\n"
if stderr:
Expand Down

0 comments on commit c3fe334

Please sign in to comment.