Make detection of command line flags that prevent sudo from overriding TERMINFO more robust

kitty_tests/shell_integration.py

@@ -32,6 +32,33 @@ def bash_ok():
     return int(major_ver) >= 5 and relstatus == 'release'
 
 
+def extract_sudo_function(
+    content: str, opening='sudo() {', closing='}',
+    witht='command sudo TERMINFO="$TERMINFO" "$@";', without='command sudo "$@";'
+) -> str:
+    """Extract the sudo() function from bash/zsh shell integration content using indentation."""
+    lines = content.split('\n')
+    start_idx = None
+    indent_len = None
+    for i, line in enumerate(lines):
+        stripped = line.lstrip()
+        if stripped == opening:
+            start_idx = i
+            indent_len = len(line) - len(stripped)
+            break
+    if start_idx is None:
+        return None
+    for i in range(start_idx + 1, len(lines)):
+        line = lines[i]
+        stripped = line.strip()
+        if stripped == closing and (len(line) - len(line.lstrip())) == indent_len:
+            func = '\n'.join(lines[start_idx:i + 1])
+            func = func.replace(witht, 'echo "with_terminfo";')
+            func = func.replace(without, 'echo "no_terminfo";')
+            return func
+    return None
+
+
 def basic_shell_env(home_dir):
     ans = {
         'PATH': os.environ.get('PATH', '/usr/local/bin:/usr/bin:/bin:/usr/sbin:/sbin'),
@@ -115,6 +142,69 @@ def run_shell(self, shell='zsh', rc='', cmd='', setup_env=None, extra_env=None):
                         continue
                     raise
 
+    @unittest.skipUnless(shutil.which('zsh'), 'zsh not installed')
+    def test_zsh_sudo_parser(self):
+        if self.with_kitten:
+            return
+        src = os.path.join(shell_integration_dir, 'zsh', 'kitty-integration')
+        with open(src) as f:
+            func = extract_sudo_function(f.read())
+        self.assertIsNotNone(func)
+        self.sudo_parser_tests(['zsh', '--no-rcs', '-c'], func)
+
+    @unittest.skipUnless(bash_ok(), 'bash not installed, too old, or debug build')
+    def test_bash_sudo_parser(self):
+        if self.with_kitten:
+            return
+        src = os.path.join(shell_integration_dir, 'bash', 'kitty.bash')
+        with open(src) as f:
+            func = extract_sudo_function(f.read())
+        self.assertIsNotNone(func)
+        self.sudo_parser_tests(['bash', '--noprofile', '--norc', '-c'], func)
+
+    @unittest.skipUnless(shutil.which('fish'), 'fish not installed')
+    def test_fish_sudo_parser(self):
+        if self.with_kitten:
+            return
+        src = os.path.join(shell_integration_dir, 'fish', 'vendor_conf.d', 'kitty-shell-integration.fish')
+        with open(src) as f:
+            func = extract_sudo_function(
+                f.read(), opening='function sudo', closing='end', witht='sudo TERMINFO="$TERMINFO" $argv',
+                without='sudo $argv')
+        self.assertIsNotNone(func)
+        self.sudo_parser_tests(['fish', '--no-config', '-c'], func)
+
+    def sudo_parser_tests(self, shell: list[str], func: str) -> None:
+        test_cases = [
+            (['ls'],                               'with_terminfo', 'simple command'),
+            (['-e', 'file.txt'],                   'no_terminfo',   '-e flag'),
+            (['--edit', 'file.txt'],               'no_terminfo',   '--edit flag'),
+            (['-v'],                               'no_terminfo',   '-v validate flag'),
+            (['--validate'],                       'no_terminfo',   '--validate flag'),
+            (['-nv'],                              'no_terminfo',   '-nv grouped flags'),
+            (['-ne', 'file.txt'],                  'no_terminfo',   '-ne grouped flags'),
+            (['-u', 'root', 'ls'],                 'with_terminfo', '-u takes argument'),
+            (['-u', 'root', '-v'],                 'no_terminfo',   '-u arg then -v'),
+            (['--', '-v'],                         'with_terminfo', '-- ends option processing'),
+            (['TERM=xterm-kitty', 'ls'],           'with_terminfo', 'env var before command'),
+            (['TERM=xterm-kitty', '-v', 'ls'],     'no_terminfo',   'env var before -v'),
+            (['-g', 'wheel', '-v'],                'no_terminfo', '-g takes argument'),
+            (['-D', '/tmp', '-e'],                 'no_terminfo', '-D takes argument'),
+            (['-R', '/chroot', '-e'],              'no_terminfo', '-R takes argument'),
+            (['-h', 'localhost', '-v'],            'no_terminfo', '-h takes argument'),
+            (['-uroot', '-v', 'ls'],               'with_terminfo', '-u with embedded argument'),
+            (['--user', '-v', 'ls'],               'with_terminfo', '--user takes argument'),
+            (['--user=root', '-v', 'ls'],          'no_terminfo', '--user=root self-contained'),
+        ]
+        for cmd, expected, _ in test_cases:
+            func += f'\n\nsudo {" ".join(cmd)}'
+        cp = subprocess.run(shell + [func], capture_output=True)
+        self.assertEqual(cp.returncode, 0, f'{shell[0]} failed with stderr: {cp.stderr.decode()}')
+        for i, line in enumerate(cp.stdout.decode().splitlines()):
+            line = line.strip()
+            cmd, expected, desc = test_cases[i]
+            self.assertEqual(expected, line, str(cmd))
+
     @unittest.skipUnless(shutil.which('zsh'), 'zsh not installed')
     def test_zsh_integration(self):
         ps1, rps1 = 'left>', '<right'

shell-integration/bash/kitty.bash

@@ -18,23 +18,33 @@ if [[ -n "$KITTY_BASH_INJECT" ]]; then
 
     if [[ -n "$ksi_val" && "$ksi_val" != *no-sudo* && -n "$TERMINFO" && ! ( -r "/usr/share/terminfo/x/xterm-kitty" || -r "/usr/share/terminfo/78/xterm-kitty" ) ]]; then
         # this must be done before sourcing user bashrc otherwise aliasing of sudo does not work
-        sudo() {
-            # Ensure terminfo is available in sudo
-            builtin local is_sudoedit="n"
-            for arg; do
-                if [[ "$arg" == "-e" || $arg == "--edit" ]]; then
-                    is_sudoedit="y"
-                    builtin break;
+            sudo() {
+                # Ensure terminfo is available in sudo
+                builtin local cannot_set_env_var="n"
+                builtin local ignore_arg="n"
+                for arg; do
+                    if [[ "$ignore_arg" == "y" ]]; then ignore_arg="n"; builtin continue; fi
+                    case "$arg" in
+                        --) builtin break ;;  # end of options
+                        # sudo -e disallows setting env vars
+                        --edit | -e* | -[!-]*e*) cannot_set_env_var="y"; builtin break ;;
+                        # sudo-rs -v disallows setting env vars
+                        --validate | -v* | -[!-]*v*) cannot_set_env_var="y"; builtin break ;;
+                        # flags that require a following argument
+                        --user|--group|--host|--chdir|--chroot|--role|--type|--command-timeout|--auth-type|--login-class|--prompt|--close-from|--other-user) ignore_arg="y" ;;
+                        --*) ;;  # misc long opt
+                        -*a|-*C|-*c|-*D|-*g|-*h|-*p|-*R|-*r|-*t|-*T|-*u) ignore_arg="y" ;;
+                        # flag or env var setting
+                        -* | *=*) ;;
+                        *) builtin break ;;  # command found
+                    esac
+                done
+                if [[ "$cannot_set_env_var" == "y" ]]; then
+                    builtin command sudo "$@";
+                else
+                    builtin command sudo TERMINFO="$TERMINFO" "$@";
                 fi
-                [[ "$arg" == "--" ]] && builtin break  # end of options
-                [[ "$arg" != -* && "$arg" != *=* ]] && builtin break  # command found
-            done
-            if [[ "$is_sudoedit" == "y" ]]; then
-                builtin command sudo "$@";
-            else
-                builtin command sudo TERMINFO="$TERMINFO" "$@";
-            fi
-        }
+            }
     fi
 
     if [[ "$kitty_bash_inject" == *"posix"* ]]; then

shell-integration/fish/vendor_conf.d/kitty-shell-integration.fish

@@ -144,20 +144,42 @@ function __ksi_schedule --on-event fish_prompt -d "Setup kitty integration after
         and not test -r "/usr/share/terminfo/x/xterm-kitty" -o -r "/usr/share/terminfo/78/xterm-kitty"
         # Ensure terminfo is available in sudo
         function sudo
-            set --local is_sudoedit "n"
+            set --local cannot_set_env_var "n"
+            set --local ignore_arg "n"
             for arg in $argv
-                if string match -q -- "-e" "$arg" or string match -q -- "--edit" "$arg"
-                    set is_sudoedit "y"
-                    break
+                if string match -q -- "$ignore_arg" "y"
+                    set ignore_arg "n"
+                    continue
                 end
-                if string match -q -- "--" "$arg"
+                if string match -r -q -- '^--$' "$arg"
                     break  # end of options
                 end
-                if not string match -r -q -- "^-" "$arg" and not string match -r -q -- "=" "$arg"
-                    break  # reached the command
+                if string match -r -q -- '^-[^-]*e' "$arg" or string match -q -- "--edit" "$arg"
+                    set cannot_set_env_var "y"  # sudo -e
+                    break
                 end
+                if string match -r -q -- '^-[^-]*v' "$arg" or string match -q -- "--validate" "$arg"
+                    set cannot_set_env_var "y"  # sudo-rs -e
+                    break
+                end
+                if string match -r -q -- '--user|--group|--host|--chdir|--chroot|--role|--type|--command-timeout|--auth-type|--login-class|--prompt|--close-from|--other-user' "$arg"
+
+                    set ignore_arg="y"
+                    continue
+                end
+                if string match -r -q -- '--.+' "$arg"
+                    continue
+                end
+                if string match -r -q -- '-.*a|-.*C|-.*c|-.*D|-.*g|-.*h|-.*p|-.*R|-.*r|-.*t|-.*T|-.*u)' "$arg"
+                    set ignore_arg="y"
+                    continue
+                end
+                if string match -r -q -- '-.+' "$arg" or string match -r -q -- '.+=.+' "$arg"
+                    continue
+                end
+                break  # command found
             end
-            if string match -q -- "$is_sudoedit" "y"
+            if string match -q -- "$cannot_set_env_var" "y"
                 command sudo $argv
             else
                 command sudo TERMINFO="$TERMINFO" $argv

shell-integration/zsh/kitty-integration

@@ -406,16 +406,26 @@ _ksi_deferred_init() {
         if [[ -n "$TERMINFO" && ! ( -r "/usr/share/terminfo/x/xterm-kitty" || -r "/usr/share/terminfo/78/xterm-kitty" ) ]]; then
             sudo() {
                 # Ensure terminfo is available in sudo
-                builtin local is_sudoedit="n"
+                builtin local cannot_set_env_var="n"
+                builtin local ignore_arg="n"
                 for arg; do
-                    if [[ "$arg" == "-e" || $arg == "--edit" ]]; then
-                        is_sudoedit="y"
-                        builtin break;
-                    fi
-                    [[ "$arg" == "--" ]] && builtin break  # end of options
-                    [[ "$arg" != -* && "$arg" != *=* ]] && builtin break  # command found
+                    if [[ "$ignore_arg" == "y" ]]; then ignore_arg="n"; builtin continue; fi
+                    case "$arg" in
+                        --) builtin break ;;  # end of options
+                        # sudo -e disallows setting env vars
+                        --edit | -e* | -[!-]*e*) cannot_set_env_var="y"; builtin break ;;
+                        # sudo-rs -v disallows setting env vars
+                        --validate | -v* | -[!-]*v*) cannot_set_env_var="y"; builtin break ;;
+                        # flags that require a following argument
+                        --user|--group|--host|--chdir|--chroot|--role|--type|--command-timeout|--auth-type|--login-class|--prompt|--close-from|--other-user) ignore_arg="y" ;;
+                        --*) ;;  # misc long opt
+                        -*a|-*C|-*c|-*D|-*g|-*h|-*p|-*R|-*r|-*t|-*T|-*u) ignore_arg="y" ;;
+                        # flag or env var setting
+                        -* | *=*) ;;
+                        *) builtin break ;;  # command found
+                    esac
                 done
-                if [[ "$is_sudoedit" == "y" ]]; then
+                if [[ "$cannot_set_env_var" == "y" ]]; then
                     builtin command sudo "$@";
                 else
                     builtin command sudo TERMINFO="$TERMINFO" "$@";