From 68c77de3471bcd6a79921a5253d169d9447b5476 Mon Sep 17 00:00:00 2001
From: Maas Lalani <maas@lalani.dev>
Date: Sun, 19 Nov 2023 15:34:00 -0500
Subject: [PATCH] feat: implement detection of kitty keyboard protocol

---
 output.go       |  10 ++--
 termenv_unix.go | 127 +++++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 131 insertions(+), 6 deletions(-)

diff --git a/output.go b/output.go
index b8e82df..259779b 100644
--- a/output.go
+++ b/output.go
@@ -215,11 +215,11 @@ func (o Output) WriteString(s string) (int, error) {
 //
 // The byte returned represents the bitset of supported flags.
 //
-//   0b1     (01) — Disambiguate escape codes
-//   0b10    (02) — Report event types
-//   0b100   (04) — Report alternate keys
-//   0b1000  (08) — Report all keys as escape codes
-//   0b10000 (16) — Report associated text
+// 0b1     (01) — Disambiguate escape codes
+// 0b10    (02) — Report event types
+// 0b100   (04) — Report alternate keys
+// 0b1000  (08) — Report all keys as escape codes
+// 0b10000 (16) — Report associated text
 //
 func (o Output) KittyKeyboardProtocolSupport() byte {
 	f := func() {
diff --git a/termenv_unix.go b/termenv_unix.go
index 8648277..91c1c82 100644
--- a/termenv_unix.go
+++ b/termenv_unix.go
@@ -114,7 +114,76 @@ func (o Output) backgroundColor() Color {
 }
 
 func (o Output) kittyKeyboardProtocolSupport() byte {
-	return 0b11111
+	// screen/tmux can't support OSC, because they can be connected to multiple
+	// terminals concurrently.
+	term := o.environ.Getenv("TERM")
+	if strings.HasPrefix(term, "screen") || strings.HasPrefix(term, "tmux") {
+		return 0b00000
+	}
+
+	tty := o.TTY()
+	if tty == nil {
+		return 0b00000
+	}
+
+	if !o.unsafe {
+		fd := int(tty.Fd())
+		// if in background, we can't control the terminal
+		if !isForeground(fd) {
+			return 0b00000
+		}
+
+		t, err := unix.IoctlGetTermios(fd, tcgetattr)
+		if err != nil {
+			return 0b00000
+		}
+		defer unix.IoctlSetTermios(fd, tcsetattr, t) //nolint:errcheck
+
+		noecho := *t
+		noecho.Lflag = noecho.Lflag &^ unix.ECHO
+		noecho.Lflag = noecho.Lflag &^ unix.ICANON
+		if err := unix.IoctlSetTermios(fd, tcsetattr, &noecho); err != nil {
+			return 0b00000
+		}
+	}
+
+	// first, send CSI query to see whether this terminal supports the
+	// kitty keyboard protocol
+	fmt.Fprintf(tty, CSI+"?u")
+
+	// then, query primary device data, should be supported by all terminals
+	// if we receive a response for the primary device data befor the kitty keyboard
+	// protocol response, this terminal does not support kitty keyboard protocol.
+	fmt.Fprintf(tty, CSI+"c")
+
+	response, isAttrs, err := o.readNextResponseKittyKeyboardProtocol()
+
+	// we queried for the kitty keyboard protocol current progressive enhancements
+	// but received the primary device attributes response, therefore this terminal
+	// does not support the kitty keyboard protocol.
+	if err != nil || isAttrs {
+		return 0
+	}
+
+	// read the primary attrs response and ignore it.
+	_, _, err = o.readNextResponseKittyKeyboardProtocol()
+	if err != nil {
+		return 0
+	}
+
+	// we receive a valid response to the kitty keyboard protocol query, this
+	// terminal supports the protocol.
+	//
+	// parse the response and return the flags supported.
+	//
+	//   0    1 2 3 4
+	//   \x1b [ ? 1 u
+	//
+	if len(response) <= 3 {
+		return 0
+	}
+
+	return response[3]
 }
 
 func (o *Output) waitForData(timeout time.Duration) error {
@@ -161,6 +230,62 @@ func (o *Output) readNextByte() (byte, error) {
 	return b[0], nil
 }
 
+// readNextResponseKittyKeyboardProtocol reads either a CSI response to the current
+// progressive enhancement status or primary device attributes response.
+//   - CSI response: "\x1b]?31u"
+//   - primary device attributes response: "\x1b]?64;1;2;7;8;9;15;18;21;44;45;46c"
+func (o *Output) readNextResponseKittyKeyboardProtocol() (response string, isAttrs bool, err error) {
+	start, err := o.readNextByte()
+	if err != nil {
+		return "", false, ErrStatusReport
+	}
+
+	// first byte must be ESC
+	for start != ESC {
+		start, err = o.readNextByte()
+		if err != nil {
+			return "", false, ErrStatusReport
+		}
+	}
+
+	response += string(start)
+
+	// next byte is [
+	tpe, err := o.readNextByte()
+	if err != nil {
+		return "", false, ErrStatusReport
+	}
+	response += string(tpe)
+
+	if tpe != '[' {
+		return "", false, ErrStatusReport
+	}
+
+	for {
+		b, err := o.readNextByte()
+		if err != nil {
+			return "", false, ErrStatusReport
+		}
+		response += string(b)
+
+		switch b {
+		case 'u':
+			// kitty keyboard protocol response
+			return response, false, nil
+		case 'c':
+			// primary device attributes response
+			return response, true, nil
+		}
+
+		// both responses have less than 38 bytes, so if we read more, that's an error
+		if len(response) > 38 {
+			break
+		}
+	}
+
+	return response, isAttrs, nil
+}
+
 // readNextResponse reads either an OSC response or a cursor position response:
 //   - OSC response: "\x1b]11;rgb:1111/1111/1111\x1b\\"
 //   - cursor position response: "\x1b[42;1R"