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

This commit is contained in:
Kovid Goyal 2024-09-29 20:36:12 +05:30
parent 4bb0d3dbfb
commit af83d855de
No known key found for this signature in database
GPG key ID: 06BC317B515ACE7C
10 changed files with 219 additions and 29 deletions

View file

@ -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]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

View file

@ -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
--------------------

View file

@ -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))

View file

@ -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

View file

@ -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:

View file

@ -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:

View file

@ -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.

View file

@ -453,6 +453,8 @@ def launch_child(
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 @@ def launch_child(
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 @@ def new_window(
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 @@ def new_special_window(
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 @@ def new_special_window(
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')

View file

@ -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)

View file

@ -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 {