Merge pull request #14 from Comp-490-TR-Afternoon-Group-Four/#7-confirm-reboot-env-variable

#7 confirm reboot env variable

README.md

@@ -456,6 +456,7 @@ Or via environment variables:
 * `THEFUCK_RULES` – list of enabled rules, like `DEFAULT_RULES:rm_root` or `sudo:no_command`;
 * `THEFUCK_EXCLUDE_RULES` – list of disabled rules, like `git_pull:git_push`;
 * `THEFUCK_REQUIRE_CONFIRMATION` – require confirmation before running new command, `true/false`;
+* `THEFUCK_REQUIRE_DOUBLE_CONFIRMATION` – require double confirmation before running potentially volitile commands (eg. `reboot`), value of `true/false`;
 * `THEFUCK_WAIT_COMMAND` – max amount of time in seconds for getting previous command output;
 * `THEFUCK_NO_COLORS` – disable colored output, `true/false`;
 * `THEFUCK_PRIORITY` – priority of the rules, like `no_command=9999:apt_get=100`,

tests/test_ui.py

@@ -12,6 +12,7 @@ def patch_get_key(monkeypatch):
     def patch(vals):
         vals = iter(vals)
         monkeypatch.setattr('thefuck.ui.get_key', lambda: next(vals))
+        monkeypatch.setenv("THEFUCK_REQUIRE_DOUBLE_CONFIRMATION", 'True')
 
     return patch
 
@@ -62,6 +63,10 @@ class TestSelectCommand(object):
         return [CorrectedCommand('ls', None, 100),
                 CorrectedCommand('cd', None, 100)]
 
+    @pytest.fixture
+    def reboot_command(self):
+        return [CorrectedCommand('reboot', None, 100)]
+
     def test_without_commands(self, capsys):
         assert ui.select_command(iter([])) is None
         assert capsys.readouterr() == ('', 'No fucks given\n')
@@ -106,3 +111,12 @@ class TestSelectCommand(object):
             u'{mark}\x1b[1K\rcd [enter/↑/↓/ctrl+c]\n'
         ).format(mark=const.USER_COMMAND_MARK)
         assert capsys.readouterr() == ('', stderr)
+
+    def test_with_double_confirmation(self, capsys, patch_get_key, reboot_command):
+        patch_get_key(['\n'])
+        assert ui.select_command(iter(reboot_command)) == reboot_command[0]
+
+    def test_with_double_confirmation_abort(self, capsys, patch_get_key, reboot_command):
+        patch_get_key([const.KEY_CTRL_C])
+        assert ui.select_command(iter(reboot_command)) is None
+        assert capsys.readouterr() == ('', const.USER_COMMAND_MARK + u'\x1b[1K\rreboot [enter/↑/↓/ctrl+c]\nAborted\n')

thefuck/const.py

@@ -28,10 +28,13 @@ ALL_ENABLED = _GenConst('All rules enabled')
 DEFAULT_RULES = [ALL_ENABLED]
 DEFAULT_PRIORITY = 1000
 
+DOUBLE_CONFIRMATION_SCRIPTS = {"reboot": "Are you sure you would like to reboot the system?"}
+
 DEFAULT_SETTINGS = {'rules': DEFAULT_RULES,
                     'exclude_rules': [],
                     'wait_command': 3,
                     'require_confirmation': True,
+                    'require_double_confirmation': False,
                     'no_colors': False,
                     'debug': False,
                     'priority': {},
@@ -49,6 +52,7 @@ ENV_TO_ATTR = {'THEFUCK_RULES': 'rules',
                'THEFUCK_EXCLUDE_RULES': 'exclude_rules',
                'THEFUCK_WAIT_COMMAND': 'wait_command',
                'THEFUCK_REQUIRE_CONFIRMATION': 'require_confirmation',
+               'THEFUCK_REQUIRE_DOUBLE_CONFIRMATION': 'require_double_confirmation',
                'THEFUCK_NO_COLORS': 'no_colors',
                'THEFUCK_DEBUG': 'debug',
                'THEFUCK_PRIORITY': 'priority',

thefuck/logs.py

@@ -72,6 +72,21 @@ def confirm_text(corrected_command):
             blue=color(colorama.Fore.BLUE)))
 
 
+def double_confirm_text(confirmation_text):
+    sys.stderr.write(
+        (u'{prefix}{clear}{bold}{text}{reset} '
+         u'[{green}enter{reset}'
+         u'/{red}ctrl+c{reset}]').format(
+            prefix=const.USER_COMMAND_MARK,
+            text=confirmation_text,
+            clear='\033[1K\r',
+            bold=color(colorama.Style.BRIGHT),
+            green=color(colorama.Fore.GREEN),
+            red=color(colorama.Fore.RED),
+            reset=color(colorama.Style.RESET_ALL),
+        ))
+
+
 def debug(msg):
     if settings.debug:
         sys.stderr.write(u'{blue}{bold}DEBUG:{reset} {msg}\n'.format(

thefuck/ui.py

@@ -80,10 +80,13 @@ def select_command(corrected_commands):
 
     logs.confirm_text(selector.value)
 
-    for action in read_actions():
+    selected = False
+    while not selected:
+        action = next(read_actions())
+
         if action == const.ACTION_SELECT:
             sys.stderr.write('\n')
-            return selector.value
+            selected = True
         elif action == const.ACTION_ABORT:
             logs.failed('\nAborted')
             return
@@ -93,3 +96,21 @@ def select_command(corrected_commands):
         elif action == const.ACTION_NEXT:
             selector.next()
             logs.confirm_text(selector.value)
+
+    if settings.require_double_confirmation and selector.value.script in const.DOUBLE_CONFIRMATION_SCRIPTS:
+        return double_confirm(selector)
+
+    return selector.value
+
+
+def double_confirm(selector):
+    confirmation_text = const.DOUBLE_CONFIRMATION_SCRIPTS[selector.value.script]
+    logs.double_confirm_text(confirmation_text)
+
+    for action in read_actions():
+        if action == const.ACTION_SELECT:
+            sys.stderr.write('\n')
+            return selector.value
+        elif action == const.ACTION_ABORT:
+            logs.failed('\nAborted')
+            return