From c5bbdd557f70d58bd8123d97127da0db4860501c Mon Sep 17 00:00:00 2001
From: Zac Bergquist <zac.bergquist@goteleport.com>
Date: Sun, 22 Sep 2024 15:54:47 -0600
Subject: [PATCH] tsh play --format=text

For sessions with print events, this command will strip ANSI
control sequences from the recording and write the text directly
to standard out. This makes it easier to grep for a particular
pattern or just dump a session recording to a file for manual
analysis.

Closes #11694
---
 go.mod                  |  1 +
 go.sum                  |  2 ++
 tool/tsh/common/play.go | 33 +++++++++++++++++++++++++++++++--
 tool/tsh/common/tsh.go  |  4 ++--
 4 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/go.mod b/go.mod
index d845750718e9b..c34584934a941 100644
--- a/go.mod
+++ b/go.mod
@@ -157,6 +157,7 @@ require (
 	github.com/opensearch-project/opensearch-go/v2 v2.3.0
 	github.com/parquet-go/parquet-go v0.23.0
 	github.com/pavlo-v-chernykh/keystore-go/v4 v4.5.0
+	github.com/pborman/ansi v1.0.0
 	github.com/pelletier/go-toml v1.9.5
 	github.com/pkg/sftp v1.13.6
 	github.com/pquerna/otp v1.4.0
diff --git a/go.sum b/go.sum
index f83eb553e710d..54292184ce9dd 100644
--- a/go.sum
+++ b/go.sum
@@ -1995,6 +1995,8 @@ github.com/paulmach/orb v0.11.1/go.mod h1:5mULz1xQfs3bmQm63QEJA6lNGujuRafwA5S/En
 github.com/paulmach/protoscan v0.2.1/go.mod h1:SpcSwydNLrxUGSDvXvO0P7g7AuhJ7lcKfDlhJCDw2gY=
 github.com/pavlo-v-chernykh/keystore-go/v4 v4.5.0 h1:2nosf3P75OZv2/ZO/9Px5ZgZ5gbKrzA3joN1QMfOGMQ=
 github.com/pavlo-v-chernykh/keystore-go/v4 v4.5.0/go.mod h1:lAVhWwbNaveeJmxrxuSTxMgKpF6DjnuVpn6T8WiBwYQ=
+github.com/pborman/ansi v1.0.0 h1:OqjHMhvlSuCCV5JT07yqPuJPQzQl+WXsiZ14gZsqOrQ=
+github.com/pborman/ansi v1.0.0/go.mod h1:SgWzwMAx1X/Ez7i90VqF8LRiQtx52pWDiQP+x3iGnzw=
 github.com/pborman/uuid v1.2.1 h1:+ZZIw58t/ozdjRaXh/3awHfmWRbzYxJoAdNJxe/3pvw=
 github.com/pborman/uuid v1.2.1/go.mod h1:X/NO0urCmaxf9VXbdlT7C2Yzkj2IKimNn4k+gtPdI/k=
 github.com/pelletier/go-toml v1.9.5 h1:4yBQzkHv+7BHq2PQUZF3Mx0IYxG7LsP222s7Agd3ve8=
diff --git a/tool/tsh/common/play.go b/tool/tsh/common/play.go
index 14900ce485871..bddb56feb0fe2 100644
--- a/tool/tsh/common/play.go
+++ b/tool/tsh/common/play.go
@@ -30,6 +30,7 @@ import (
 
 	"github.com/ghodss/yaml"
 	"github.com/gravitational/trace"
+	"github.com/pborman/ansi"
 
 	"github.com/gravitational/teleport"
 	apievents "github.com/gravitational/teleport/api/types/events"
@@ -115,9 +116,10 @@ func exportSession(cf *CLIConf) error {
 	}
 
 	switch format {
-	case teleport.JSON, teleport.YAML:
+	case teleport.JSON, teleport.YAML, teleport.Text:
 	default:
-		return trace.Errorf("Invalid format %s, only json and yaml are supported", format)
+		// this should be unreachable since kingpin validates the format flag
+		return trace.BadParameter("Invalid format %s", format)
 	}
 
 	sid, err := session.ParseID(cf.SessionID)
@@ -144,6 +146,8 @@ func exportSession(cf *CLIConf) error {
 		exporter = jsonSessionExporter{}
 	case teleport.YAML:
 		exporter = yamlSessionExporter{}
+	case teleport.Text:
+		exporter = textSessionExporter{}
 	}
 
 	exporter.WriteStart()
@@ -233,6 +237,31 @@ func (yamlSessionExporter) WriteEvent(evt apievents.AuditEvent) error {
 	return err
 }
 
+type textSessionExporter struct{}
+
+func (textSessionExporter) WriteStart() error     { return nil }
+func (textSessionExporter) WriteEnd() error       { return nil }
+func (textSessionExporter) WriteSeparator() error { return nil }
+
+func (textSessionExporter) WriteEvent(evt apievents.AuditEvent) error {
+	printEvent, ok := evt.(*apievents.SessionPrint)
+	if !ok {
+		return nil
+	}
+
+	var writeErr error
+
+	// strip ANSI control characters
+	if stripped, err := ansi.Strip(printEvent.Data); err != nil {
+		// best we can do is just write the original data
+		_, writeErr = os.Stdout.Write(printEvent.Data)
+	} else {
+		_, writeErr = os.Stdout.Write(stripped)
+	}
+
+	return writeErr
+}
+
 // exportFile converts the binary protobuf events from the file
 // identified by path to text (JSON/YAML) and writes the converted
 // events to standard out.
diff --git a/tool/tsh/common/tsh.go b/tool/tsh/common/tsh.go
index 500dcab53c79f..79bf18f491934 100644
--- a/tool/tsh/common/tsh.go
+++ b/tool/tsh/common/tsh.go
@@ -983,8 +983,8 @@ func Run(ctx context.Context, args []string, opts ...CliOption) error {
 	play.Flag("speed", "Playback speed, applicable when streaming SSH or Kubernetes sessions.").Default("1x").EnumVar(&cf.PlaySpeed, "0.5x", "1x", "2x", "4x", "8x")
 	play.Flag("skip-idle-time", "Quickly skip over idle time, applicable when streaming SSH or Kubernetes sessions.").BoolVar(&cf.NoWait)
 	play.Flag("format", defaults.FormatFlagDescription(
-		teleport.PTY, teleport.JSON, teleport.YAML,
-	)).Short('f').Default(teleport.PTY).EnumVar(&cf.Format, teleport.PTY, teleport.JSON, teleport.YAML)
+		teleport.PTY, teleport.JSON, teleport.YAML, teleport.Text,
+	)).Short('f').Default(teleport.PTY).EnumVar(&cf.Format, teleport.PTY, teleport.JSON, teleport.YAML, teleport.Text)
 	play.Arg("session-id", "ID or path to session file to play").Required().StringVar(&cf.SessionID)
 
 	// scp