Add a framework for easily and securely using remote control from the main function of a custom kitten

kovidgoyal · kovidgoyal · commit af83d855de05 · 2024-09-29T20:36:12.000+05:30
diff --git a/docs/changelog.rst b/docs/changelog.rst
@@ -74,6 +74,11 @@ consumption to do the same tasks.
 Detailed list of changes
 -------------------------------------
 
+0.37.0 [future]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+- Custom kittens: Add :ref:`a framework <kitten_main_rc>` for easily and securely using remote control from within a kitten's :code:`main()` function
+
 0.36.4 [2024-09-27]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
diff --git a/docs/kittens/custom.rst b/docs/kittens/custom.rst
@@ -231,6 +231,56 @@ The function will only send the event if the program is receiving events of
 that type, and will return ``True`` if it sent the event, and ``False`` if not.
 
 
+.. _kitten_main_rc:
+
+Using remote control inside the main() kitten function
+------------------------------------------------------------
+
+You can use kitty's remote control features inside the main() function of a
+kitten, even without enabling remote control. This is useful if you want to
+probe kitty for more information before presenting some UI to the user or if
+you want the user to be able to control kitty from within your kitten's UI
+rather than after it has finished running. To enable it, simply tell kitty your kitten
+requires remote control, as shown in the example below::
+
+    import json
+    import sys
+    from pprint import pprint
+
+    from kittens.tui.handler import kitten_ui
+
+    @kitten_ui(allow_remote_control=True)
+    def main(args: list[str]) -> str:
+        # get the result of running kitten @ ls
+        cp = main.remote_control(['ls'], capture_output=True)
+        if cp.returncode != 0:
+            sys.stderr.buffer.write(cp.stderr)
+            raise SystemExit(cp.returncode)
+        output = json.loads(cp.stdout)
+        pprint(output)
+        # open a new tab with a title specified by the user
+        title = input('Enter the name of tab: ')
+        window_id = main.remote_control(['launch', '--type=tab', '--tab-title', title], check=True, capture_output=True).stdout.decode()
+        return window_id
+
+:code:`allow_remote_control=True` tells kitty to run this kitten with remote
+control enabled, regardless of whether it is enabled globally or not.
+To run a remote control command use the :code:`main.remote_control()` function
+which is a thin wrapper around Python's :code:`subprocess.run` function. Note
+that by default, for security, child processes launched by your kitten cannot use remote
+control, thus it is necessary to use :code:`main.remote_control()`. If you wish
+to enable child processes to use remote control, call
+:code:`main.allow_indiscriminate_remote_control()`.
+
+Remote control access can be further secured by using
+:code:`kitten_ui(allow_remote_control=True, remote_control_password='ls set-colors')`.
+This will use a secure generated password to restrict remote control.
+You can specify a space separated list of remote control commands to allow, see
+:opt:`remote_control_password` for details. The password value is accessible
+as :code:`main.password` and is used by :code:`main.remote_control()`
+automatically.
+
+
 Debugging kittens
 --------------------
 
diff --git a/kittens/runner.py b/kittens/runner.py
@@ -80,18 +80,22 @@ class KittenMetadata(NamedTuple):
     no_ui: bool = False
     has_ready_notification: bool = False
     open_url_handler: Optional[Callable[[BossType, WindowType, str, int, str], bool]] = None
-
+    allow_remote_control: bool = False
+    remote_control_password: str | bool = False
 
 
 def create_kitten_handler(kitten: str, orig_args: List[str]) -> KittenMetadata:
     from kitty.constants import config_dir
     kitten = resolved_kitten(kitten)
     m = import_kitten_main_module(config_dir, kitten)
+    main = m['start']
     handle_result = m['end']
     return KittenMetadata(
         handle_result=partial(handle_result, [kitten] + orig_args),
         type_of_input=getattr(handle_result, 'type_of_input', None),
         no_ui=getattr(handle_result, 'no_ui', False),
+        allow_remote_control=getattr(main, 'allow_remote_control', False),
+        remote_control_password=getattr(main, 'remote_control_password', True),
         has_ready_notification=getattr(handle_result, 'has_ready_notification', False),
         open_url_handler=getattr(handle_result, 'open_url_handler', None))
 
diff --git a/kittens/tui/handler.py b/kittens/tui/handler.py
@@ -2,12 +2,14 @@
 # License: GPL v3 Copyright: 2018, Kovid Goyal <kovid at kovidgoyal.net>
 
 
+import os
 from collections import deque
 from contextlib import suppress
 from types import TracebackType
 from typing import TYPE_CHECKING, Any, Callable, ContextManager, Deque, Dict, NamedTuple, Optional, Sequence, Type, Union, cast
 
-from kitty.fast_data_types import monotonic
+from kitty.constants import kitten_exe, running_in_kitty
+from kitty.fast_data_types import monotonic, safe_pipe
 from kitty.types import DecoratedFunc, ParsedShortcut
 from kitty.typing import (
     AbstractEventLoop,
@@ -47,6 +49,98 @@ def is_click(a: ButtonEvent, b: ButtonEvent) -> bool:
     return x*x + y*y <= 4
 
 
+class KittenUI:
+    allow_remote_control: bool = False
+    remote_control_password: bool | str = False
+
+    def __init__(self, func: Callable[[list[str]], str], allow_remote_control: bool, remote_control_password: bool | str):
+        self.func = func
+        self.allow_remote_control = allow_remote_control
+        self.remote_control_password = remote_control_password
+        self.password = self.to = ''
+        self.rc_fd = -1
+        self.initialized = False
+
+    def initialize(self) -> None:
+        if self.initialized:
+            return
+        self.initialized = True
+        if running_in_kitty():
+            return
+        if self.allow_remote_control:
+            self.to = os.environ.get('KITTY_LISTEN_ON', '')
+            self.rc_fd = int(self.to.partition(':')[-1])
+            os.set_inheritable(self.rc_fd, False)
+        if (self.remote_control_password or self.remote_control_password == '') and not self.password:
+            import socket
+            with socket.fromfd(self.rc_fd, socket.AF_UNIX, socket.SOCK_STREAM) as s:
+                data = s.recv(256)
+            if not data.endswith(b'\n'):
+                raise Exception(f'The remote control password was invalid: {data!r}')
+            self.password = data.strip().decode()
+
+    def __call__(self, args: list[str]) -> str:
+        self.initialize()
+        return self.func(args)
+
+    def allow_indiscriminate_remote_control(self, enable: bool = True) -> None:
+        if self.rc_fd > -1:
+            if enable:
+                os.set_inheritable(self.rc_fd, True)
+                if self.password:
+                    os.environ['KITTY_RC_PASSWORD'] = self.password
+            else:
+                os.set_inheritable(self.rc_fd, False)
+                if self.password:
+                    os.environ.pop('KITTY_RC_PASSWORD', None)
+
+    def remote_control(self, cmd: str | Sequence[str], **kw: Any) -> Any:
+        if not self.allow_remote_control:
+            raise ValueError('Remote control is not enabled, remember to use allow_remote_control=True')
+        prefix = [kitten_exe(), '@']
+        r = -1
+        pass_fds = list(kw.get('pass_fds') or ())
+        try:
+            if self.rc_fd > -1:
+                pass_fds.append(self.rc_fd)
+            if self.password and self.rc_fd > -1:
+                r, w = safe_pipe(False)
+                os.write(w, self.password.encode())
+                os.close(w)
+                prefix += ['--password-file', f'fd:{r}', '--use-password', 'always']
+                pass_fds.append(r)
+            if pass_fds:
+                kw['pass_fds'] = tuple(pass_fds)
+            if isinstance(cmd, str):
+                cmd = ' '.join(prefix)
+            else:
+                cmd = prefix + list(cmd)
+            import subprocess
+            if self.rc_fd > -1:
+                is_inheritable = os.get_inheritable(self.rc_fd)
+                if not is_inheritable:
+                    os.set_inheritable(self.rc_fd, True)
+            try:
+                return subprocess.run(cmd, **kw)
+            finally:
+                if self.rc_fd > -1 and not is_inheritable:
+                    os.set_inheritable(self.rc_fd, False)
+        finally:
+            if r > -1:
+                os.close(r)
+
+
+def kitten_ui(
+    allow_remote_control: bool = KittenUI.allow_remote_control,
+    remote_control_password: bool | str = KittenUI.allow_remote_control,
+) -> Callable[[Callable[[list[str]], str]], KittenUI]:
+
+    def wrapper(impl: Callable[..., Any]) -> KittenUI:
+        return KittenUI(impl, allow_remote_control, remote_control_password)
+
+    return wrapper
+
+
 class Handler:
 
     image_manager_class: Optional[Type[ImageManagerType]] = None
diff --git a/kitty/boss.py b/kitty/boss.py
@@ -1948,17 +1948,32 @@ def run_kitten_with_metadata(
             else:
                 cmd = [kitty_exe(), '+runpy', 'from kittens.runner import main; main()']
                 env['PYTHONWARNINGS'] = 'ignore'
-            overlay_window = tab.new_special_window(
-                SpecialWindow(
-                    cmd + final_args,
-                    stdin=data,
-                    env=env,
-                    cwd=w.cwd_of_child,
-                    overlay_for=w.id,
-                    overlay_behind=end_kitten.has_ready_notification,
-                ),
-                copy_colors_from=w
-            )
+            remote_control_fd = -1
+            if end_kitten.allow_remote_control:
+                remote_control_passwords: Optional[dict[str, Sequence[str]]] = None
+                initial_data = b''
+                if end_kitten.remote_control_password:
+                    from secrets import token_hex
+                    p = token_hex(16)
+                    remote_control_passwords = {p: end_kitten.remote_control_password if isinstance(end_kitten.remote_control_password, str) else ''}
+                    initial_data = p.encode() + b'\n'
+                remote = self.add_fd_based_remote_control(remote_control_passwords, initial_data)
+                remote_control_fd = remote.fileno()
+            try:
+                overlay_window = tab.new_special_window(
+                    SpecialWindow(
+                        cmd + final_args,
+                        stdin=data,
+                        env=env,
+                        cwd=w.cwd_of_child,
+                        overlay_for=w.id,
+                        overlay_behind=end_kitten.has_ready_notification,
+                    ),
+                    copy_colors_from=w, remote_control_fd=remote_control_fd,
+                )
+            finally:
+                if end_kitten.allow_remote_control:
+                    remote.close()
             wid = w.id
             overlay_window.actions_on_close.append(partial(self.on_kitten_finish, wid, custom_callback or end_kitten.handle_result, default_data=default_data))
             overlay_window.open_url_handler = end_kitten.open_url_handler
@@ -2351,9 +2366,11 @@ def special_window_for_cmd(
         overlay_for = w.id if w and as_overlay else None
         return SpecialWindow(cmd, input_data, cwd_from=cwd_from, overlay_for=overlay_for, env=env)
 
-    def add_fd_based_remote_control(self, remote_control_passwords: Optional[dict[str, Sequence[str]]] = None) -> socket.socket:
+    def add_fd_based_remote_control(self, remote_control_passwords: Optional[dict[str, Sequence[str]]] = None, initial_data: bytes = b'') -> socket.socket:
         local, remote = socket.socketpair()
         os.set_inheritable(remote.fileno(), True)
+        if initial_data:
+            local.send(initial_data)
         lfd = os.dup(local.fileno())
         local.close()
         try:
diff --git a/kitty/child.py b/kitty/child.py
@@ -7,7 +7,7 @@
 from collections.abc import Generator, Sequence
 from contextlib import contextmanager, suppress
 from itertools import count
-from typing import TYPE_CHECKING, DefaultDict, Optional, Protocol, Union
+from typing import TYPE_CHECKING, DefaultDict, Optional
 
 import kitty.fast_data_types as fast_data_types
 
@@ -23,11 +23,6 @@
     from .window import CwdRequest
 
 
-class InheritableFile(Protocol):
-
-    def fileno(self) -> int: ...
-
-
 if is_macos:
     from kitty.fast_data_types import cmdline_of_process as cmdline_
     from kitty.fast_data_types import cwd_of_process as _cwd
@@ -216,13 +211,15 @@ def __init__(
         is_clone_launch: str = '',
         add_listen_on_env_var: bool = True,
         hold: bool = False,
-        pass_fds: tuple[Union[int, InheritableFile], ...] = (),
+        pass_fds: tuple[int, ...] = (),
+        remote_control_fd: int = -1,
     ):
         self.is_clone_launch = is_clone_launch
         self.id = next(child_counter)
         self.add_listen_on_env_var = add_listen_on_env_var
         self.argv = list(argv)
         self.pass_fds = pass_fds
+        self.remote_control_fd = remote_control_fd
         if cwd_from:
             try:
                 cwd = cwd_from.modify_argv_for_launch_with_cwd(self.argv, env) or cwd
@@ -251,7 +248,9 @@ def get_final_env(self) -> dict[str, str]:
         env['COLORTERM'] = 'truecolor'
         env['KITTY_PID'] = getpid()
         env['KITTY_PUBLIC_KEY'] = boss.encryption_public_key
-        if self.add_listen_on_env_var and boss.listening_on:
+        if self.remote_control_fd > -1:
+            env['KITTY_LISTEN_ON'] = f'fd:{self.remote_control_fd}'
+        elif self.add_listen_on_env_var and boss.listening_on:
             env['KITTY_LISTEN_ON'] = boss.listening_on
         else:
             env.pop('KITTY_LISTEN_ON', None)
@@ -299,6 +298,9 @@ def fork(self) -> Optional[int]:
         self.final_env = self.get_final_env()
         argv = list(self.argv)
         cwd = self.cwd
+        pass_fds = self.pass_fds
+        if self.remote_control_fd > -1:
+            pass_fds += self.remote_control_fd,
         if self.should_run_via_run_shell_kitten:
             # bash will only source ~/.bash_profile if it detects it is a login
             # shell (see the invocation section of the bash man page), which it
@@ -319,7 +321,7 @@ def fork(self) -> Optional[int]:
             if ksi == 'invalid':
                 ksi = 'enabled'
             argv = [kitten_exe(), 'run-shell', '--shell', shlex.join(argv), '--shell-integration', ksi]
-            if is_macos and not self.pass_fds and not opts.forward_stdio:
+            if is_macos and not pass_fds and not opts.forward_stdio:
                 # In addition for getlogin() to work we need to run the shell
                 # via the /usr/bin/login wrapper, sigh.
                 # And login on macOS looks for .hushlogin in CWD instead of
@@ -339,12 +341,10 @@ def fork(self) -> Optional[int]:
             argv = cmdline_for_hold(argv)
             final_exe = argv[0]
         env = tuple(f'{k}={v}' for k, v in self.final_env.items())
-        pass_fds = tuple(sorted(x if isinstance(x, int) else x.fileno() for x in self.pass_fds))
         pid = fast_data_types.spawn(
             final_exe, cwd, tuple(argv), env, master, slave, stdin_read_fd, stdin_write_fd,
             ready_read_fd, ready_write_fd, tuple(handled_signals), kitten_exe(), opts.forward_stdio, pass_fds)
         os.close(slave)
-        self.pass_fds = ()
         self.pid = pid
         self.child_fd = master
         if stdin is not None:
diff --git a/kitty/remote_control.py b/kitty/remote_control.py
@@ -279,6 +279,7 @@ def handle_cmd(
 default=rc-pass
 A file from which to read the password. Trailing whitespace is ignored. Relative
 paths are resolved from the kitty configuration directory. Use - to read from STDIN.
+Use :code:`fd:num` to read from the file descriptor :code:`num`.
 Used if no :option:`kitten @ --password` is supplied. Defaults to checking for the
 :file:`rc-pass` file in the kitty configuration directory.
 
diff --git a/kitty/tabs.py b/kitty/tabs.py
diff --git a/tools/cmd/at/main.go b/tools/cmd/at/main.go
diff --git a/tools/cmd/at/socket_io.go b/tools/cmd/at/socket_io.go
