kitty/tools/cmd/edit_in_kitty/main.go

// License: GPLv3 Copyright: 2022, Kovid Goyal, <kovid at kovidgoyal.net>

package edit_in_kitty

import (
	"bytes"
	"fmt"
	"io"
	"io/fs"
	"os"
	"strconv"
	"strings"

	"github.com/emmansun/base64"
	"golang.org/x/sys/unix"

	"github.com/kovidgoyal/kitty/tools/cli"
	"github.com/kovidgoyal/kitty/tools/tui"
	"github.com/kovidgoyal/kitty/tools/tui/loop"
	"github.com/kovidgoyal/kitty/tools/utils"
	"github.com/kovidgoyal/kitty/tools/utils/humanize"
)

var _ = fmt.Print

func encode(x string) string {
	return base64.StdEncoding.EncodeToString(utils.UnsafeStringToBytes(x))
}

type OnDataCallback = func(data_type string, data []byte) error

func edit_loop(data_to_send string, kill_if_signaled bool, on_data OnDataCallback) (exit_code int, err error) {
	lp, err := loop.New(loop.NoAlternateScreen, loop.NoRestoreColors, loop.NoMouseTracking)
	if err != nil {
		return
	}
	current_text := strings.Builder{}
	data := strings.Builder{}
	data.Grow(4096)
	started := false
	canceled := false
	update_type := ""

	handle_line := func(line string) error {
		if canceled {
			return nil
		}
		if started {
			if update_type == "" {
				update_type = line
			} else {
				if line == "KITTY_DATA_END" {
					lp.QueueWriteString(update_type + "\r\n")
					if update_type == "DONE" {
						// Parse the editor exit code from the data payload.
						// Decoding errors are intentionally ignored: if the data is
						// missing or malformed (e.g. from an older kitty server),
						// we fall back to exit_code=0 for backward compatibility.
						if data.Len() > 0 {
							if b, err2 := base64.StdEncoding.DecodeString(data.String()); err2 == nil {
								if n, err3 := strconv.Atoi(strings.TrimSpace(string(b))); err3 == nil {
									exit_code = n
								}
							}
						}
						lp.Quit(exit_code)
						return nil
					}
					b, err := base64.StdEncoding.DecodeString(data.String())
					data.Reset()
					data.Grow(4096)
					started = false
					if err == nil {
						err = on_data(update_type, b)
					}
					update_type = ""
					if err != nil {
						return err
					}
				} else {
					data.WriteString(line)
				}
			}
		} else {
			if line == "KITTY_DATA_START" {
				started = true
				update_type = ""
			}
		}
		return nil
	}

	check_for_line := func() error {
		if canceled {
			return nil
		}
		s := current_text.String()
		for {
			idx := strings.Index(s, "\n")
			if idx < 0 {
				break
			}
			err = handle_line(s[:idx])
			if err != nil {
				return err
			}
			s = s[idx+1:]
		}
		current_text.Reset()
		current_text.Grow(4096)
		if s != "" {
			current_text.WriteString(s)
		}
		return nil
	}

	lp.OnInitialize = func() (string, error) {
		pos, chunk_num := 0, 0
		for {
			limit := min(pos+2048, len(data_to_send))
			if limit <= pos {
				break
			}
			lp.QueueWriteString("\x1bP@kitty-edit|" + strconv.Itoa(chunk_num) + ":")
			lp.QueueWriteString(data_to_send[pos:limit])
			lp.QueueWriteString("\x1b\\")
			chunk_num++
			pos = limit
		}
		lp.QueueWriteString("\x1bP@kitty-edit|\x1b\\")
		return "", nil
	}

	lp.OnText = func(text string, from_key_event bool, in_bracketed_paste bool) error {
		if !from_key_event {
			current_text.WriteString(text)
			err = check_for_line()
			if err != nil {
				return err
			}
		}
		return nil
	}

	const abort_msg = "\x1bP@kitty-edit|0:abort_signaled=interrupt\x1b\\\x1bP@kitty-edit|\x1b\\"

	lp.OnKeyEvent = func(event *loop.KeyEvent) error {
		if event.MatchesPressOrRepeat("ctrl+c") || event.MatchesPressOrRepeat("esc") {
			event.Handled = true
			canceled = true
			lp.QueueWriteString(abort_msg)
			if !started {
				return tui.Canceled
			}
		}
		return nil
	}

	err = lp.Run()
	if err != nil {
		return
	}
	if canceled {
		return 1, tui.Canceled
	}

	ds := lp.DeathSignalName()
	if ds != "" {
		fmt.Print(abort_msg)
		if kill_if_signaled {
			lp.KillIfSignalled()
			return
		}
		return 1, &tui.KilledBySignal{Msg: fmt.Sprint("Killed by signal: ", ds), SignalName: ds}
	}
	return
}

func edit_in_kitty(path string, opts *Options) (exit_code int, err error) {
	read_file, err := os.Open(path)
	if err != nil {
		return 1, fmt.Errorf("Failed to open %s for reading with error: %w", path, err)
	}
	defer read_file.Close()
	var s unix.Stat_t
	err = unix.Fstat(int(read_file.Fd()), &s)
	if err != nil {
		return 1, fmt.Errorf("Failed to stat %s with error: %w", path, err)
	}
	if s.Size > int64(opts.MaxFileSize)*1024*1024 {
		return 1, fmt.Errorf("File size %s is too large for performant editing", humanize.Bytes(uint64(s.Size)))
	}

	file_data, err := io.ReadAll(read_file)
	if err != nil {
		return 1, fmt.Errorf("Failed to read from %s with error: %w", path, err)
	}
	read_file.Close()
	data := strings.Builder{}
	data.Grow(len(file_data) * 4)

	add := func(key, val string) {
		if data.Len() > 0 {
			data.WriteString(",")
		}
		data.WriteString(key)
		data.WriteString("=")
		data.WriteString(val)
	}
	add_encoded := func(key, val string) { add(key, encode(val)) }

	if unix.Access(path, unix.R_OK|unix.W_OK) != nil {
		return 1, fmt.Errorf("%s is not readable and writeable", path)
	}
	cwd, err := os.Getwd()
	if err != nil {
		return 1, fmt.Errorf("Failed to get the current working directory with error: %w", err)
	}
	add_encoded("cwd", cwd)
	for _, arg := range os.Args[2:] {
		add_encoded("a", arg)
	}
	add("file_inode", fmt.Sprintf("%d:%d:%d", s.Dev, s.Ino, s.Mtim.Nano()))
	add_encoded("file_data", utils.UnsafeBytesToString(file_data))
	fmt.Println("Waiting for editing to be completed, press Esc to abort...")
	write_data := func(data_type string, rdata []byte) (err error) {
		err = utils.AtomicWriteFile(path, bytes.NewReader(rdata), fs.FileMode(s.Mode).Perm())
		if err != nil {
			err = fmt.Errorf("Failed to write data to %s with error: %w", path, err)
		}
		return
	}
	exit_code, err = edit_loop(data.String(), true, write_data)
	if err != nil {
		if err == tui.Canceled {
			return 1, err
		}
		return 1, fmt.Errorf("Failed to receive edited file back from terminal with error: %w", err)
	}
	return
}

type Options struct {
	MaxFileSize int
}

func EntryPoint(parent *cli.Command) *cli.Command {
	sc := parent.AddSubCommand(&cli.Command{
		Name:             "edit-in-kitty",
		Usage:            "[options] [+lnum] file-to-edit",
		ShortDescription: "Edit a file in a kitty overlay window",
		HelpText: "Edit the specified file in a kitty overlay window. Works over SSH as well.\n\n" +
			"For usage instructions see: https://sw.kovidgoyal.net/kitty/shell-integration/#edit-file",
		Run: func(cmd *cli.Command, args []string) (ret int, err error) {
			if len(args) == 0 {
				fmt.Fprintln(os.Stderr, "Usage:", cmd.Usage)
				return 1, fmt.Errorf("No file to edit specified.")
			}

			var file_path string
			if len(args) == 1 {
				file_path = args[0]
			} else if len(args) == 2 && strings.HasPrefix(args[0], "+") {
				var lnum string
				lnum, file_path = args[0][1:], args[1]
				if _, err := strconv.Atoi(lnum); err != nil {
					return 1, fmt.Errorf("Invalid line number %s", lnum)
				}
			} else {
				fmt.Fprintln(os.Stderr, "Usage:", cmd.Usage)
				return 1, fmt.Errorf("Only one file to edit and optionally a line number must be specified")
			}

			var opts Options
			err = cmd.GetOptionValues(&opts)
			if err != nil {
				return 1, err
			}
			return edit_in_kitty(file_path, &opts)
		},
	})
	AddCloneSafeOpts(sc)
	sc.Add(cli.OptionSpec{
		Name:    "--max-file-size",
		Default: "8",
		Type:    "int",
		Help:    "The maximum allowed size (in MB) of files to edit. Since the file data has to be base64 encoded and transmitted over the tty device, overly large files will not perform well.",
	})
	return sc
}