From b4cc38a1d9e2694c5d1ce0f30659be3422a97de4 Mon Sep 17 00:00:00 2001
From: Kovid Goyal <kovid@kovidgoyal.net>
Date: Mon, 28 Feb 2022 12:59:26 +0530
Subject: [PATCH] Fix symlink extraction over ssh

---
 kittens/ssh/copy.py                |  7 +++--
 kitty_tests/ssh.py                 | 44 +++++++++++++++++++++++++++++-
 shell-integration/ssh/bootstrap.sh |  1 +
 3 files changed, 48 insertions(+), 4 deletions(-)

diff --git a/kittens/ssh/copy.py b/kittens/ssh/copy.py
index 781788a2e..7048af3cc 100644
--- a/kittens/ssh/copy.py
+++ b/kittens/ssh/copy.py
@@ -12,6 +12,7 @@ from typing import (
 from kitty.cli import parse_args
 from kitty.cli_stub import CopyCLIOptions
 from kitty.types import run_once
+from ..transfer.utils import expand_home, home_path
 
 
 @run_once
@@ -48,9 +49,9 @@ def parse_copy_args(args: Optional[Sequence[str]] = None) -> Tuple[CopyCLIOption
 
 
 def resolve_file_spec(spec: str, is_glob: bool) -> Iterator[str]:
-    ans = os.path.expandvars(os.path.expanduser(spec))
+    ans = os.path.expandvars(expand_home(spec))
     if not os.path.isabs(ans):
-        ans = os.path.expanduser(f'~/{ans}')
+        ans = expand_home(f'~/{ans}')
     if is_glob:
         files = glob.glob(ans)
         if not files:
@@ -93,7 +94,7 @@ def parse_copy_instructions(val: str) -> Iterable[Tuple[str, CopyInstruction]]:
         raise CopyCLIError('No files to copy specified')
     if len(locations) > 1 and opts.dest:
         raise CopyCLIError('Specifying a remote location with more than one file is not supported')
-    home = os.path.expanduser('~')
+    home = home_path()
     for loc in locations:
         arcname = get_arcname(loc, opts.dest, home)
         yield loc, CopyInstruction(arcname, tuple(opts.exclude))
diff --git a/kitty_tests/ssh.py b/kitty_tests/ssh.py
index cc79c9666..c265ea2f3 100644
--- a/kitty_tests/ssh.py
+++ b/kitty_tests/ssh.py
@@ -11,6 +11,7 @@ from functools import lru_cache
 from kittens.ssh.config import load_config, options_for_host
 from kittens.ssh.main import bootstrap_script, get_connection_data
 from kittens.ssh.options.utils import DELETE_ENV_VAR
+from kittens.transfer.utils import set_paths
 from kitty.constants import is_macos
 from kitty.fast_data_types import CURSOR_BEAM
 from kitty.options.utils import shell_integration
@@ -20,6 +21,12 @@ from . import BaseTest
 from .shell_integration import bash_ok, basic_shell_env
 
 
+def files_in(path):
+    for record in os.walk(path):
+        for f in record[-1]:
+            yield os.path.relpath(os.path.join(record[0], f), path)
+
+
 class SSHKitten(BaseTest):
 
     def test_basic_pty_operations(self):
@@ -70,7 +77,42 @@ print(' '.join(map(str, buf)))'''), lines=13, cols=77)
         return tuple(sh for sh in ('dash', 'zsh', 'bash', 'posh', 'sh') if shutil.which(sh))
 
     def test_ssh_copy(self):
-        pass
+        simple_data = 'rkjlhfwf9whoaa'
+
+        def touch(p):
+            with open(os.path.join(local_home, p), 'w') as f:
+                f.write(simple_data)
+
+        for sh in self.all_possible_sh:
+            with self.subTest(sh=sh), tempfile.TemporaryDirectory() as remote_home, tempfile.TemporaryDirectory() as local_home, set_paths(home=local_home):
+                tuple(map(touch, 'simple-file g.1 g.2'.split()))
+                os.makedirs(f'{local_home}/d1/d2/d3')
+                touch('d1/d2/x')
+                os.symlink('d2/x', f'{local_home}/d1/y')
+
+                conf = '''\
+copy simple-file
+copy --glob g.*
+copy d1
+'''
+                copy = load_config(overrides=filter(None, conf.splitlines()))['*'].copy
+                self.check_bootstrap(
+                    sh, remote_home, extra_exec='env; exit 0', SHELL_INTEGRATION_VALUE='',
+                    ssh_opts={'copy': copy}
+                )
+                self.assertTrue(os.path.lexists(f'{remote_home}/.terminfo/78'))
+                self.assertTrue(os.path.exists(f'{remote_home}/.terminfo/78/xterm-kitty'))
+                self.assertTrue(os.path.exists(f'{remote_home}/.terminfo/x/xterm-kitty'))
+                with open(os.path.join(remote_home, 'simple-file'), 'r') as f:
+                    self.ae(f.read(), simple_data)
+                self.assertTrue(os.path.lexists(f'{remote_home}/d1/y'))
+                self.assertTrue(os.path.exists(f'{remote_home}/d1/y'))
+                self.ae(os.readlink(f'{remote_home}/d1/y'), 'd2/x')
+                contents = set(files_in(remote_home))
+                contents.discard('.zshrc')  # added by check_bootstrap()
+                self.ae(contents, {
+                    'g.1', 'g.2', '.terminfo/kitty.terminfo', 'simple-file', '.terminfo/x/xterm-kitty', 'd1/d2/x', 'd1/y',
+                })
 
     def test_ssh_env_vars(self):
         for sh in self.all_possible_sh:
diff --git a/shell-integration/ssh/bootstrap.sh b/shell-integration/ssh/bootstrap.sh
index cb7e8c94e..c92ef9d25 100644
--- a/shell-integration/ssh/bootstrap.sh
+++ b/shell-integration/ssh/bootstrap.sh
@@ -41,6 +41,7 @@ mv_files_and_dirs() {
     cwd="$PWD";
     cd "$1";
     command find . -type d -exec mkdir -p "$2/{}" ";"
+    command find . -type l -exec sh -c "tgt=\$(command readlink -n \"{}\"); command ln -s \"\$tgt\" \"$2/{}\"; command rm -f \"{}\"" ";"
     command find . -type f -exec mv "{}" "$2/{}" ";"
     cd "$cwd";
 }