diff --git a/promptshell/__init__.py b/promptshell/__init__.py index a68927d..d1f2e39 100644 --- a/promptshell/__init__.py +++ b/promptshell/__init__.py @@ -1 +1 @@ -__version__ = "0.1.0" \ No newline at end of file +__version__ = "0.1.1" \ No newline at end of file diff --git a/promptshell/ai_terminal_assistant.py b/promptshell/ai_terminal_assistant.py index ee381be..9839465 100644 --- a/promptshell/ai_terminal_assistant.py +++ b/promptshell/ai_terminal_assistant.py @@ -231,10 +231,23 @@ def execute_command(self, user_input: str) -> str: choice = questionary.confirm(f"Do you want to run the command '{command}'?").ask() if choice: if command.startswith("CONFIRM:"): - confirmation = questionary.confirm(f"Warning: This command may be destructive. Are you sure you want to run '{command[8:]}'?").ask() + dangerous_command = command[8:].strip() + + confirmation = questionary.confirm( + f"Warning: This command may be destructive. Are you sure you want to run '{dangerous_command}'?" + ).ask() + if not confirmation: return format_text('red') + "Command execution aborted." + reset_format() - command = command[8:] + + # New safeguard: ask user to retype command + manual_input = questionary.text( + "Please manually type or paste the command to confirm:").ask() + + if manual_input.strip() != dangerous_command: + return format_text('red') + "Command mismatch. Execution aborted." + reset_format() + + command = dangerous_command formatted_command = format_text('cyan') + f"Command: {command}" + reset_format() print(formatted_command) self.command_history.append(command) diff --git a/tests/test_destructive.py b/tests/test_destructive.py new file mode 100644 index 0000000..5e7447d --- /dev/null +++ b/tests/test_destructive.py @@ -0,0 +1,57 @@ +import pytest +import os +from unittest.mock import patch, MagicMock +from promptshell.ai_terminal_assistant import AITerminalAssistant + +@pytest.fixture +def assistant(): + return AITerminalAssistant(model_name="test-model") + +@patch('promptshell.ai_terminal_assistant.questionary.confirm') +@patch('promptshell.ai_terminal_assistant.questionary.text') +@patch('os.get_terminal_size', return_value=os.terminal_size((80, 24))) +def test_confirm_safeguard_success(mock_terminal, mock_text, mock_confirm, assistant, capsys): + user_input = "delete all files" + generated_command = "CONFIRM:rm -rf *" + + assistant.command_executor = MagicMock(return_value=generated_command) + + mock_confirm.return_value.ask.return_value = True + mock_text.return_value.ask.return_value = "rm -rf *" + + assistant.execute_command(user_input) + captured = capsys.readouterr() + + assert "Command: rm -rf *" in captured.out + + +@patch('promptshell.ai_terminal_assistant.questionary.confirm') +@patch('promptshell.ai_terminal_assistant.questionary.text') +@patch('os.get_terminal_size', return_value=os.terminal_size((80, 24))) +def test_confirm_safeguard_mismatch(mock_terminal, mock_text, mock_confirm, assistant): + user_input = "delete all files" + generated_command = "CONFIRM:rm -rf *" + + assistant.command_executor = MagicMock(return_value=generated_command) + + mock_confirm.return_value.ask.return_value = True + mock_text.return_value.ask.return_value = "rm -rf /wrong/path" + + result = assistant.execute_command(user_input) + + assert "Command mismatch" in result + + +@patch('promptshell.ai_terminal_assistant.questionary.confirm') +@patch('os.get_terminal_size', return_value=os.terminal_size((80, 24))) +def test_confirm_safeguard_first_cancel(mock_terminal, mock_confirm, assistant, capsys): + user_input = "delete all files" + generated_command = "CONFIRM:rm -rf *" + + assistant.command_executor = MagicMock(return_value=generated_command) + mock_confirm.return_value.ask.return_value = False + + assistant.execute_command(user_input) + captured = capsys.readouterr() + + assert "Command cancelled!" in captured.out