From af83d855de0524b77b5e01c5827bb672815d9aa3 Mon Sep 17 00:00:00 2001 From: Kovid Goyal Date: Sun, 29 Sep 2024 20:36:12 +0530 Subject: [PATCH] Add a framework for easily and securely using remote control from the main function of a custom kitten --- docs/changelog.rst | 5 ++ docs/kittens/custom.rst | 50 ++++++++++++++++++++ kittens/runner.py | 6 ++- kittens/tui/handler.py | 96 ++++++++++++++++++++++++++++++++++++++- kitty/boss.py | 41 ++++++++++++----- kitty/child.py | 22 ++++----- kitty/remote_control.py | 1 + kitty/tabs.py | 14 ++++-- tools/cmd/at/main.go | 11 +++++ tools/cmd/at/socket_io.go | 2 +- 10 files changed, 219 insertions(+), 29 deletions(-) diff --git a/docs/changelog.rst b/docs/changelog.rst index 790db0ced..10e8b4d0e 100644 --- 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 ` 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 index 023dce1f5..d66740e32 100644 --- 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 index b7d3846da..16537573a 100644 --- 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 index e3c5ed24e..0f4ee5522 100644 --- a/kittens/tui/handler.py +++ b/kittens/tui/handler.py @@ -2,12 +2,14 @@ # License: GPL v3 Copyright: 2018, Kovid Goyal +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 index 8c54d7d8f..246426abb 100644 --- a/kitty/boss.py +++ b/kitty/boss.py @@ -1948,17 +1948,32 @@ class Boss: 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 @@ class Boss: 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 index 57cae8a5c..c87bda2ed 100644 --- a/kitty/child.py +++ b/kitty/child.py @@ -7,7 +7,7 @@ from collections import defaultdict 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 @@ if TYPE_CHECKING: 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 @@ class Child: 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 @@ class Child: 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 @@ class Child: 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 @@ class Child: 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 @@ class Child: 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 index 5537ed925..fae3d02b1 100644 --- a/kitty/remote_control.py +++ b/kitty/remote_control.py @@ -279,6 +279,7 @@ completion=type:file relative:conf kwds:- 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 index f8386bc0e..ae673035f 100644 --- a/kitty/tabs.py +++ b/kitty/tabs.py @@ -453,6 +453,8 @@ class Tab: # {{{ is_clone_launch: str = '', add_listen_on_env_var: bool = True, hold: bool = False, + pass_fds: tuple[int, ...] = (), + remote_control_fd: int = -1, ) -> Child: check_for_suitability = True if cmd is None: @@ -500,7 +502,9 @@ class Tab: # {{{ pwid = platform_window_id(self.os_window_id) if pwid is not None: fenv['WINDOWID'] = str(pwid) - ans = Child(cmd, cwd or self.cwd, stdin, fenv, cwd_from, is_clone_launch=is_clone_launch, add_listen_on_env_var=add_listen_on_env_var, hold=hold) + ans = Child( + cmd, cwd or self.cwd, stdin, fenv, cwd_from, is_clone_launch=is_clone_launch, + add_listen_on_env_var=add_listen_on_env_var, hold=hold, pass_fds=pass_fds, remote_control_fd=remote_control_fd) ans.fork() return ans @@ -532,11 +536,13 @@ class Tab: # {{{ remote_control_passwords: Optional[dict[str, Sequence[str]]] = None, hold: bool = False, bias: Optional[float] = None, + pass_fds: tuple[int, ...] = (), + remote_control_fd: int = -1, ) -> Window: child = self.launch_child( use_shell=use_shell, cmd=cmd, stdin=stdin, cwd_from=cwd_from, cwd=cwd, env=env, is_clone_launch=is_clone_launch, add_listen_on_env_var=False if allow_remote_control and remote_control_passwords else True, - hold=hold, + hold=hold, pass_fds=pass_fds, remote_control_fd=remote_control_fd, ) window = Window( self, child, self.args, override_title=override_title, @@ -561,6 +567,8 @@ class Tab: # {{{ copy_colors_from: Optional[Window] = None, allow_remote_control: bool = False, remote_control_passwords: Optional[dict[str, Sequence[str]]] = None, + pass_fds: tuple[int, ...] = (), + remote_control_fd: int = -1, ) -> Window: return self.new_window( use_shell=False, cmd=special_window.cmd, stdin=special_window.stdin, @@ -568,7 +576,7 @@ class Tab: # {{{ cwd_from=special_window.cwd_from, cwd=special_window.cwd, overlay_for=special_window.overlay_for, env=special_window.env, location=location, copy_colors_from=copy_colors_from, allow_remote_control=allow_remote_control, watchers=special_window.watchers, overlay_behind=special_window.overlay_behind, - hold=special_window.hold, remote_control_passwords=remote_control_passwords, + hold=special_window.hold, remote_control_passwords=remote_control_passwords, pass_fds=pass_fds, remote_control_fd=remote_control_fd, ) @ac('win', 'Close all windows in the tab other than the currently active window') diff --git a/tools/cmd/at/main.go b/tools/cmd/at/main.go index 826807f7d..acea9449c 100644 --- a/tools/cmd/at/main.go +++ b/tools/cmd/at/main.go @@ -326,6 +326,17 @@ func get_password(password string, password_file string, password_env string, us ttyf.Close() } } + } else if strings.HasPrefix(password_file, "fd:") { + var fd int + if fd, err = strconv.Atoi(password_file[3:]); err == nil { + f := os.NewFile(uintptr(fd), password_file) + var q []byte + if q, err = io.ReadAll(f); err == nil { + ans.is_set = true + ans.val = string(q) + } + f.Close() + } } else { var q []byte q, err = os.ReadFile(password_file) diff --git a/tools/cmd/at/socket_io.go b/tools/cmd/at/socket_io.go index ff74d4940..8ca8e8353 100644 --- a/tools/cmd/at/socket_io.go +++ b/tools/cmd/at/socket_io.go @@ -170,7 +170,7 @@ func do_socket_io(io_data *rc_io_data) (serialized_response []byte, err error) { f := os.NewFile(uintptr(fd), "fd:"+global_options.to_address) conn, err = net.FileConn(f) if err != nil { - return nil, err + return nil, fmt.Errorf("Failed to open a socket for the remote control file descriptor: %d with error: %w", fd, err) } defer f.Close() } else {