mirrors/kitty
1
0
Fork 0
mirror of https://github.com/kovidgoyal/kitty synced 2026-07-06 08:01:58 +02:00
Code Issues Packages Projects Releases Wiki Activity

Allow using a --kitten argument with the ssh kitten to make it easier to override settings from ssh.conf

Browse Source
This commit is contained in:
Kovid Goyal
2022-03-06 10:07:46 +05:30
parent 31ea5d74a7
commit c713dc0ca8
6 changed files with 97 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
docs/kittens/ssh.rst
View File

@@ -50,13 +50,19 @@ quick example:
copy --dest=foo/bar some-file
copy --glob some/files.*
# Include some config from environment variables.
# This will read config directives from the contents
# of all environemnt variables starting with SSH_KITTEN_
envinclude SSH_KITTEN_*
See below for full details on the syntax and options of :file:`ssh.conf`.
Additionally, you can pass config options on the command line:
.. code-block:: sh
kitty +kitten ssh --kitten hostname=somehost --kitten interpreter=python ...
The :code:`--kitten` argument can be specified multiple times, with directives
from :file:`ssh.conf`. These are merged with :file:`ssh.conf` as if they were
appended to the end of that file. It is important to specify the hostname as
the first setting, otherwise the last hostname from :file:`ssh.conf` will be
used.
.. note::

4
kittens/ssh/config.py
View File

@@ -86,7 +86,7 @@ def load_config(*paths: str, overrides: Optional[Iterable[str]] = None) -> Dict[
return ans
def init_config() -> Dict[str, SSHOptions]:
def init_config(overrides: Optional[Iterable[str]] = None) -> Dict[str, SSHOptions]:
config = tuple(resolve_config(SYSTEM_CONF, defconf))
opts_dict = load_config(*config)
opts_dict = load_config(*config, overrides=overrides)
return opts_dict

101
kittens/ssh/main.py
View File

@@ -23,11 +23,11 @@ from typing import (
from kitty.constants import cache_dir, shell_integration_dir, terminfo_dir
from kitty.fast_data_types import get_options
from kitty.short_uuid import uuid4
from kitty.types import run_once
from kitty.utils import SSHConnectionData
from .completion import complete, ssh_options
from .config import options_for_host
from .config import init_config, options_for_host
from .copy import CopyInstruction
from .options.types import Options as SSHOptions
from .options.utils import DELETE_ENV_VAR
@@ -40,7 +40,7 @@ def serialize_env(env: Dict[str, str], base_env: Dict[str, str]) -> bytes:
for k in sorted(env):
v = env[k]
if v is DELETE_ENV_VAR:
if v == DELETE_ENV_VAR:
lines.append(f'unset {shlex.quote(k)}')
elif v == '_kitty_copy_env_var_':
q = base_env.get(k)
@@ -119,18 +119,9 @@ def make_tarfile(ssh_opts: SSHOptions, base_env: Dict[str, str]) -> bytes:
return buf.getvalue()
@run_once
def load_ssh_options() -> Dict[str, SSHOptions]:
from .config import init_config
return init_config()
def get_ssh_data(msg: str, ssh_opts: Optional[Dict[str, SSHOptions]] = None) -> Iterator[bytes]:
def get_ssh_data(msg: str) -> Iterator[bytes]:
record_sep = b'\036'
if ssh_opts is None:
ssh_opts = load_ssh_options()
def fmt_prefix(msg: Any) -> bytes:
return str(msg).encode('ascii') + record_sep
@@ -156,7 +147,9 @@ def get_ssh_data(msg: str, ssh_opts: Optional[Dict[str, SSHOptions]] = None) ->
traceback.print_exc()
yield fmt_prefix('!incorrect ssh data password')
else:
ssh_opts = {k: SSHOptions(v) for k, v in env_data['opts'].items()}
resolved_ssh_opts = options_for_host(hostname, username, ssh_opts)
resolved_ssh_opts.copy = {k: CopyInstruction(*v) for k, v in resolved_ssh_opts.copy.items()}
try:
data = make_tarfile(resolved_ssh_opts, env_data['env'])
except Exception:
@@ -175,13 +168,6 @@ def safe_remove(x: str) -> None:
def prepare_script(ans: str, replacements: Dict[str, str]) -> str:
pw = uuid4()
with tempfile.NamedTemporaryFile(prefix='ssh-kitten-pw-', suffix='.json', dir=cache_dir(), delete=False) as tf:
data = {'pw': pw, 'env': dict(os.environ)}
tf.write(json.dumps(data).encode('utf-8'))
atexit.register(safe_remove, tf.name)
replacements['DATA_PASSWORD'] = pw
replacements['PASSWORD_FILENAME'] = os.path.basename(tf.name)
for k in ('EXEC_CMD',):
replacements[k] = replacements.get(k, '')
@@ -191,14 +177,20 @@ def prepare_script(ans: str, replacements: Dict[str, str]) -> str:
return re.sub('|'.join(fr'\b{k}\b' for k in replacements), sub, ans)
def bootstrap_script(script_type: str = 'sh', **replacements: str) -> str:
def bootstrap_script(script_type: str = 'sh', exec_cmd: str = '', ssh_opts_dict: Dict[str, Dict[str, Any]] = {}) -> str:
with open(os.path.join(shell_integration_dir, 'ssh', f'bootstrap.{script_type}')) as f:
ans = f.read()
pw = uuid4()
with tempfile.NamedTemporaryFile(prefix='ssh-kitten-pw-', suffix='.json', dir=cache_dir(), delete=False) as tf:
data = {'pw': pw, 'env': dict(os.environ), 'opts': ssh_opts_dict}
tf.write(json.dumps(data).encode('utf-8'))
atexit.register(safe_remove, tf.name)
replacements = {'DATA_PASSWORD': pw, 'PASSWORD_FILENAME': os.path.basename(tf.name), 'EXEC_CMD': exec_cmd}
return prepare_script(ans, replacements)
def load_script(script_type: str = 'sh', exec_cmd: str = '') -> str:
return bootstrap_script(script_type, EXEC_CMD=exec_cmd)
def load_script(script_type: str = 'sh', exec_cmd: str = '', ssh_opts_dict: Dict[str, Dict[str, Any]] = {}) -> str:
return bootstrap_script(script_type, exec_cmd, ssh_opts_dict=ssh_opts_dict)
def get_ssh_cli() -> Tuple[Set[str], Set[str]]:
@@ -213,13 +205,22 @@ def get_ssh_cli() -> Tuple[Set[str], Set[str]]:
return boolean_ssh_args, other_ssh_args
def get_connection_data(args: List[str], cwd: str = '') -> Optional[SSHConnectionData]:
def is_extra_arg(arg: str, extra_args: Tuple[str, ...]) -> str:
for x in extra_args:
if arg == x or arg.startswith(f'{x}='):
return x
return ''
def get_connection_data(args: List[str], cwd: str = '', extra_args: Tuple[str, ...] = ()) -> Optional[SSHConnectionData]:
boolean_ssh_args, other_ssh_args = get_ssh_cli()
port: Optional[int] = None
expecting_port = expecting_identity = False
expecting_option_val = False
expecting_hostname = False
expecting_extra_val = ''
host_name = identity_file = found_ssh = ''
found_extra_args: List[Tuple[str, str]] = []
for i, arg in enumerate(args):
if not found_ssh:
@@ -247,6 +248,15 @@ def get_connection_data(args: List[str], cwd: str = '') -> Optional[SSHConnectio
else:
identity_file = arg[2:]
continue
if arg.startswith('--') and extra_args:
matching_ex = is_extra_arg(arg, extra_args)
if matching_ex:
if '=' in arg:
exval = arg.partition('=')[-1]
found_extra_args.append((matching_ex, exval))
continue
expecting_extra_val = matching_ex
expecting_option_val = True
continue
@@ -257,6 +267,8 @@ def get_connection_data(args: List[str], cwd: str = '') -> Optional[SSHConnectio
expecting_port = False
elif expecting_identity:
identity_file = arg
elif expecting_extra_val:
found_extra_args.append((expecting_extra_val, arg))
expecting_option_val = False
continue
@@ -270,7 +282,7 @@ def get_connection_data(args: List[str], cwd: str = '') -> Optional[SSHConnectio
if not os.path.isabs(identity_file):
identity_file = os.path.normpath(os.path.join(cwd or os.getcwd(), identity_file))
return SSHConnectionData(found_ssh, host_name, port, identity_file)
return SSHConnectionData(found_ssh, host_name, port, identity_file, tuple(found_extra_args))
class InvalidSSHArgs(ValueError):
@@ -285,7 +297,7 @@ class InvalidSSHArgs(ValueError):
os.execlp('ssh', 'ssh')
def parse_ssh_args(args: List[str]) -> Tuple[List[str], List[str], bool]:
def parse_ssh_args(args: List[str], extra_args: Tuple[str, ...] = ()) -> Tuple[List[str], List[str], bool, Tuple[str, ...]]:
boolean_ssh_args, other_ssh_args = get_ssh_cli()
passthrough_args = {f'-{x}' for x in 'Nnf'}
ssh_args = []
@@ -293,6 +305,8 @@ def parse_ssh_args(args: List[str]) -> Tuple[List[str], List[str], bool]:
expecting_option_val = False
passthrough = False
stop_option_processing = False
found_extra_args: List[str] = []
expecting_extra_val = ''
for argument in args:
if len(server_args) > 1 or stop_option_processing:
server_args.append(argument)
@@ -301,6 +315,16 @@ def parse_ssh_args(args: List[str]) -> Tuple[List[str], List[str], bool]:
if argument == '--':
stop_option_processing = True
continue
if extra_args:
matching_ex = is_extra_arg(argument, extra_args)
if matching_ex:
if '=' in argument:
exval = argument.partition('=')[-1]
found_extra_args.extend((matching_ex, exval))
else:
expecting_extra_val = matching_ex
expecting_option_val = True
continue
# could be a multi-character option
all_args = argument[1:]
for i, arg in enumerate(all_args):
@@ -321,16 +345,22 @@ def parse_ssh_args(args: List[str]) -> Tuple[List[str], List[str], bool]:
raise InvalidSSHArgs(f'unknown option -- {arg[1:]}')
continue
if expecting_option_val:
ssh_args.append(argument)
if expecting_extra_val:
found_extra_args.extend((expecting_extra_val, argument))
else:
ssh_args.append(argument)
expecting_option_val = False
continue
server_args.append(argument)
if not server_args:
raise InvalidSSHArgs()
return ssh_args, server_args, passthrough
return ssh_args, server_args, passthrough, tuple(found_extra_args)
def get_remote_command(remote_args: List[str], hostname: str = 'localhost', interpreter: str = 'sh') -> List[str]:
def get_remote_command(
remote_args: List[str], hostname: str = 'localhost', interpreter: str = 'sh',
ssh_opts_dict: Dict[str, Dict[str, Any]] = {}
) -> List[str]:
command_to_execute = ''
is_python = 'python' in interpreter.lower()
if remote_args:
@@ -342,7 +372,7 @@ def get_remote_command(remote_args: List[str], hostname: str = 'localhost', inte
else:
args = [c.replace("'", """'"'"'""") for c in remote_args]
command_to_execute = "exec \"$login_shell\" -c '{}'".format(' '.join(args))
sh_script = load_script(script_type='py' if is_python else 'sh', exec_cmd=command_to_execute)
sh_script = load_script(script_type='py' if is_python else 'sh', exec_cmd=command_to_execute, ssh_opts_dict=ssh_opts_dict)
return [f'{interpreter} -c {shlex.quote(sh_script)}']
@@ -351,7 +381,7 @@ def main(args: List[str]) -> NoReturn:
if args and args[0] == 'use-python':
args = args[1:] # backwards compat from when we had a python implementation
try:
ssh_args, server_args, passthrough = parse_ssh_args(args)
ssh_args, server_args, passthrough, found_extra_args = parse_ssh_args(args, extra_args=('--kitten',))
except InvalidSSHArgs as e:
e.system_exit()
cmd = ['ssh'] + ssh_args
@@ -374,7 +404,14 @@ def main(args: List[str]) -> NoReturn:
else:
hostname_for_match = hostname
hostname_for_match = hostname.split('@', 1)[-1].split(':', 1)[0]
cmd += get_remote_command(remote_args, hostname, options_for_host(hostname_for_match, uname, load_ssh_options()).interpreter)
overrides = []
pat = re.compile(r'^([a-zA-Z0-9_]+)[ \t]*=')
for i, a in enumerate(found_extra_args):
if i % 2 == 1:
overrides.append(pat.sub(r'\1 ', a.lstrip()))
so = init_config(overrides)
sod = {k: v._asdict() for k, v in so.items()}
cmd += get_remote_command(remote_args, hostname, options_for_host(hostname_for_match, uname, so).interpreter, sod)
os.execvp('ssh', cmd)

1
kitty/utils.py
View File

@@ -820,6 +820,7 @@ class SSHConnectionData(NamedTuple):
hostname: str
port: Optional[int] = None
identity_file: str = ''
extra_args: Tuple[Tuple[str, str], ...] = ()
def get_new_os_window_size(

13
kitty_tests/__init__.py
View File

@@ -97,7 +97,7 @@ class Callbacks:
def handle_remote_ssh(self, msg):
from kittens.ssh.main import get_ssh_data
if self.pty:
for line in get_ssh_data(msg, {'*': self.pty.ssh_opts} if self.pty.ssh_opts else None):
for line in get_ssh_data(msg):
self.pty.write_to_child(line)
def handle_remote_echo(self, msg):
@@ -159,9 +159,9 @@ class BaseTest(TestCase):
s = Screen(c, lines, cols, scrollback, cell_width, cell_height, 0, c)
return s
def create_pty(self, argv, cols=80, lines=25, scrollback=100, cell_width=10, cell_height=20, options=None, cwd=None, env=None, ssh_opts=None):
def create_pty(self, argv, cols=80, lines=25, scrollback=100, cell_width=10, cell_height=20, options=None, cwd=None, env=None):
self.set_options(options)
return PTY(argv, lines, cols, scrollback, cell_width, cell_height, cwd, env, ssh_opts)
return PTY(argv, lines, cols, scrollback, cell_width, cell_height, cwd, env)
def assertEqualAttributes(self, c1, c2):
x1, y1, c1.x, c1.y = c1.x, c1.y, 0, 0
@@ -174,12 +174,7 @@ class BaseTest(TestCase):
class PTY:
def __init__(self, argv, rows=25, columns=80, scrollback=100, cell_width=10, cell_height=20, cwd=None, env=None, ssh_opts=None):
if ssh_opts:
from kittens.ssh.options.types import Options as SSHOptions
self.ssh_opts = SSHOptions(ssh_opts or {})
else:
self.ssh_opts = None
def __init__(self, argv, rows=25, columns=80, scrollback=100, cell_width=10, cell_height=20, cwd=None, env=None):
if isinstance(argv, str):
argv = shlex.split(argv)
pid, self.master_fd = fork()

18
kitty_tests/ssh.py
View File

@@ -43,16 +43,18 @@ print(' '.join(map(str, buf)))'''), lines=13, cols=77)
self.ae(pty.screen_contents(), '13 77 770 260')
def test_ssh_connection_data(self):
def t(cmdline, binary='ssh', host='main', port=None, identity_file=''):
def t(cmdline, binary='ssh', host='main', port=None, identity_file='', extra_args=()):
if identity_file:
identity_file = os.path.abspath(identity_file)
q = get_connection_data(cmdline.split())
self.ae(q, SSHConnectionData(binary, host, port, identity_file))
en = set(f'{x[0]}' for x in extra_args)
q = get_connection_data(cmdline.split(), extra_args=en)
self.ae(q, SSHConnectionData(binary, host, port, identity_file, extra_args))
t('ssh main')
t('ssh un@ip -i ident -p34', host='un@ip', port=34, identity_file='ident')
t('ssh un@ip -iident -p34', host='un@ip', port=34, identity_file='ident')
t('ssh -p 33 main', port=33)
t('ssh --kitten=one -p 12 --kitten two -ix main', identity_file='x', port=12, extra_args=(('--kitten', 'one'), ('--kitten', 'two')))
def test_ssh_config_parsing(self):
def parse(conf):
@@ -198,8 +200,11 @@ copy --exclude */w.* d1
self.check_bootstrap('sh', tdir, ok_login_shell, val)
def check_bootstrap(self, sh, home_dir, login_shell='', SHELL_INTEGRATION_VALUE='enabled', extra_exec='', pre_data='', ssh_opts=None):
ssh_opts = ssh_opts or {}
if login_shell:
ssh_opts['login_shell'] = login_shell
script = bootstrap_script(
EXEC_CMD=f'echo "UNTAR_DONE"; {extra_exec}',
exec_cmd=f'echo "UNTAR_DONE"; {extra_exec}', ssh_opts_dict={'*': ssh_opts},
)
env = basic_shell_env(home_dir)
# Avoid generating unneeded completion scripts
@@ -207,10 +212,7 @@ copy --exclude */w.* d1
# prevent newuser-install from running
open(os.path.join(home_dir, '.zshrc'), 'w').close()
options = {'shell_integration': shell_integration(SHELL_INTEGRATION_VALUE or 'disabled')}
if login_shell:
ssh_opts = ssh_opts or {}
ssh_opts['login_shell'] = login_shell
pty = self.create_pty(f'{sh} -c {shlex.quote(script)}', cwd=home_dir, env=env, options=options, ssh_opts=ssh_opts)
pty = self.create_pty(f'{sh} -c {shlex.quote(script)}', cwd=home_dir, env=env, options=options)
if pre_data:
pty.write_buf = pre_data.encode('utf-8')
del script