From dd331ca12efcb753ae3d8a6cce04d7c078601195 Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Fri, 8 Apr 2022 12:10:59 +0530
Subject: [PATCH] ssh kitten: Allow changing terminal colors when connecting

---
 docs/kittens/ssh.rst              |  2 ++
 kittens/ssh/main.py               | 35 ++++++++++++++++++++++++++++++-
 kittens/ssh/options/definition.py |  9 ++++++++
 kittens/ssh/options/parse.py      |  3 +++
 kittens/ssh/options/types.py      |  2 ++
 5 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/docs/kittens/ssh.rst b/docs/kittens/ssh.rst
index 9c94993f4..453f9c4d7 100644
--- a/docs/kittens/ssh.rst
+++ b/docs/kittens/ssh.rst
@@ -9,6 +9,8 @@ Truly convenient SSH
 
 * Make kitty itself available in the remote host :opt:`on demand <kitten-ssh.remote_kitty>`
 
+* Easily :opt:`change terminal colors <kitten-ssh.color_scheme>` when connecting to remote hosts
+
 .. versionadded:: 0.25.0
    Automatic shell integration, file transfer and reuse of connections
 
diff --git a/kittens/ssh/main.py b/kittens/ssh/main.py
index a4fdb241b..c4435c9ab 100644
--- a/kittens/ssh/main.py
+++ b/kittens/ssh/main.py
@@ -28,6 +28,7 @@ from typing import (
     Union
 )
 
+from kittens.tui.operations import restore_colors, save_colors
 from kitty.constants import (
     cache_dir, runtime_dir, shell_integration_dir, ssh_control_master_template,
     str_version, terminfo_dir
@@ -35,7 +36,7 @@ from kitty.constants import (
 from kitty.options.types import Options
 from kitty.shm import SharedMemory
 from kitty.types import run_once
-from kitty.utils import SSHConnectionData, set_echo as turn_off_echo
+from kitty.utils import SSHConnectionData, set_echo as turn_off_echo, expandvars, resolve_abs_or_config_path
 
 from .completion import complete, ssh_options
 from .config import init_config
@@ -598,6 +599,34 @@ def drain_potential_tty_garbage(p: 'subprocess.Popen[bytes]', data_request: str)
                     data += q
 
 
+def change_colors(color_scheme: str) -> bool:
+    if not color_scheme:
+        return False
+    from kittens.themes.collection import load_themes, NoCacheFound, text_as_opts
+    from kittens.themes.main import colors_as_escape_codes
+    if color_scheme.endswith('.conf'):
+        conf_file = resolve_abs_or_config_path(color_scheme)
+        try:
+            with open(conf_file) as f:
+                opts = text_as_opts(f.read())
+        except FileNotFoundError:
+            raise SystemExit(f'Failed to find the color conf file: {expandvars(conf_file)}')
+    else:
+        try:
+            themes = load_themes(-1)
+        except NoCacheFound:
+            themes = load_themes()
+        cs = expandvars(color_scheme)
+        try:
+            theme = themes[cs]
+        except KeyError:
+            raise SystemExit(f'Failed to find the color theme: {cs}')
+        opts = theme.kitty_opts
+    raw = colors_as_escape_codes(opts)
+    print(save_colors(), sep='', end=raw, flush=True)
+    return True
+
+
 def run_ssh(ssh_args: List[str], server_args: List[str], found_extra_args: Tuple[str, ...]) -> NoReturn:
     cmd = [ssh_exe()] + ssh_args
     hostname, remote_args = server_args[0], server_args[1:]
@@ -650,6 +679,7 @@ def run_ssh(ssh_args: List[str], server_args: List[str], found_extra_args: Tuple
         rcmd, replacements, shm_name = get_remote_command(
             remote_args, host_opts, hostname_for_match, uname, echo_on, request_data=need_to_request_data)
         cmd += rcmd
+        colors_changed = change_colors(host_opts.color_scheme)
         try:
             p = subprocess.Popen(cmd)
         except FileNotFoundError:
@@ -661,6 +691,9 @@ def run_ssh(ssh_args: List[str], server_args: List[str], found_extra_args: Tuple
                     raise SystemExit(p.wait())
                 except KeyboardInterrupt:
                     raise SystemExit(1)
+        finally:
+            if colors_changed:
+                print(end=restore_colors(), flush=True)
 
 
 def main(args: List[str]) -> NoReturn:
diff --git a/kittens/ssh/options/definition.py b/kittens/ssh/options/definition.py
index e3afeb769..7abdabb03 100644
--- a/kittens/ssh/options/definition.py
+++ b/kittens/ssh/options/definition.py
@@ -103,6 +103,15 @@ value are expanded. The default is empty so no changing is done, which
 usually means the home directory is used.
 ''')
 
+opt('color_scheme', '', long_text='''
+Specify a color scheme to use when connecting to the remote host.  If the
+color_scheme ends with :code:`.conf` it is assumed to be the name of a config
+file to load from the kitty config directory, otherwise it is assumed to be the
+name of a color theme to load via the themes kitten. Note that only colors
+applying to the text/background are changed, other config settings in the .conf
+files/themes are ignored.
+''')
+
 opt('remote_kitty', 'if-needed', choices=('if-needed', 'no', 'yes'), long_text='''
 Make kitty available on the remote server. Useful to run kittens such as the
 icat kitten to display images or the transfer file kitten to transfer files.
diff --git a/kittens/ssh/options/parse.py b/kittens/ssh/options/parse.py
index 4c8c8c3dd..55220d0db 100644
--- a/kittens/ssh/options/parse.py
+++ b/kittens/ssh/options/parse.py
@@ -15,6 +15,9 @@ class Parser:
 
     choices_for_askpass = frozenset(('unless-set', 'ssh', 'native'))
 
+    def color_scheme(self, val: str, ans: typing.Dict[str, typing.Any]) -> None:
+        ans['color_scheme'] = str(val)
+
     def copy(self, val: str, ans: typing.Dict[str, typing.Any]) -> None:
         for k, v in copy(val, ans["copy"]):
             ans["copy"][k] = v
diff --git a/kittens/ssh/options/types.py b/kittens/ssh/options/types.py
index ee5520f38..a1b24068b 100644
--- a/kittens/ssh/options/types.py
+++ b/kittens/ssh/options/types.py
@@ -12,6 +12,7 @@ else:
 
 option_names = (  # {{{
  'askpass',
+ 'color_scheme',
  'copy',
  'cwd',
  'env',
@@ -26,6 +27,7 @@ option_names = (  # {{{
 
 class Options:
     askpass: choices_for_askpass = 'unless-set'
+    color_scheme: str = ''
     cwd: str = ''
     hostname: str = '*'
     interpreter: str = 'sh'