support assciicast and logpath refactor (#450)

* add asciicast v2 log format support and change tty log filename from username to uniqueID

* screen-recording-format support

cmd/sshpiperd/asciicast.go

@@ -0,0 +1,161 @@
+package main
+
+import (
+	"bytes"
+	"encoding/binary"
+	"encoding/json"
+	"fmt"
+	"os"
+	"path"
+	"time"
+)
+
+const (
+	msgChannelRequest     = 98
+	msgChannelOpenConfirm = 91
+)
+
+func jsonEscape(i string) string {
+	b, err := json.Marshal(i)
+	if err != nil {
+		panic(err)
+	}
+	s := string(b)
+	return s[1 : len(s)-1]
+}
+
+func readString(buf *bytes.Reader) string {
+	var l uint32
+	err := binary.Read(buf, binary.BigEndian, &l)
+	if err != nil {
+		return ""
+	}
+	s := make([]byte, l)
+	_, err = buf.Read(s)
+	if err != nil {
+		return ""
+	}
+	return string(s)
+}
+
+type asciicastLogger struct {
+	starttime    time.Time
+	envs         map[string]string
+	initWidth    uint32
+	initHeight   uint32
+	channels     map[uint32]*os.File
+	channelIDMap map[uint32]uint32
+	recorddir    string
+	prefix       string // prefix for the output file
+}
+
+func newAsciicastLogger(recorddir string, prefix string) *asciicastLogger {
+	return &asciicastLogger{
+		envs:         make(map[string]string),
+		recorddir:    recorddir,
+		channels:     make(map[uint32]*os.File),
+		channelIDMap: make(map[uint32]uint32),
+		prefix:       prefix,
+	}
+}
+
+func (l *asciicastLogger) uphook(msg []byte) ([]byte, error) {
+	if msg[0] == msgChannelData {
+		clientChannelID := binary.BigEndian.Uint32(msg[1:5])
+
+		f, ok := l.channels[clientChannelID]
+		if ok {
+			buf := msg[9:]
+			t := time.Since(l.starttime).Seconds()
+
+			_, err := fmt.Fprintf(f, "[%v,\"o\",\"%s\"]\n", t, jsonEscape(string(buf)))
+
+			if err != nil {
+				return msg, err
+			}
+		}
+	} else if msg[0] == msgChannelOpenConfirm {
+		clientChannelID := binary.BigEndian.Uint32(msg[1:5])
+		serverChannelID := binary.BigEndian.Uint32(msg[5:9])
+		l.channelIDMap[serverChannelID] = clientChannelID
+	}
+	return msg, nil
+}
+
+func (l *asciicastLogger) downhook(msg []byte) ([]byte, error) {
+	if msg[0] == msgChannelRequest {
+		t := time.Since(l.starttime).Seconds()
+		serverChannelID := binary.BigEndian.Uint32(msg[1:5])
+		clientChannelID := l.channelIDMap[serverChannelID]
+		buf := bytes.NewReader(msg[5:])
+		reqType := readString(buf)
+
+		switch reqType {
+		case "pty-req":
+			_, _ = buf.ReadByte()
+			term := readString(buf)
+			_ = binary.Read(buf, binary.BigEndian, &l.initWidth)
+			_ = binary.Read(buf, binary.BigEndian, &l.initHeight)
+			l.envs["TERM"] = term
+		case "env":
+			_, _ = buf.ReadByte()
+			varName := readString(buf)
+			varValue := readString(buf)
+			l.envs[varName] = varValue
+		case "window-change":
+			f, ok := l.channels[clientChannelID]
+			if !ok {
+				_, _ = buf.ReadByte()
+				var width, height uint32
+				_ = binary.Read(buf, binary.BigEndian, &width)
+				_ = binary.Read(buf, binary.BigEndian, &height)
+
+				_, err := fmt.Fprintf(f, "[%v,\"r\", \"%vx%v\"]\n", t, width, height)
+				if err != nil {
+					return msg, err
+				}
+			}
+		case "shell", "exec":
+			jsonEnvs, err := json.Marshal(l.envs)
+
+			if err != nil {
+				return msg, err
+			}
+
+			f, err := os.OpenFile(
+				path.Join(l.recorddir, fmt.Sprintf("%s%s-channel-%d.cast", l.prefix, reqType, clientChannelID)),
+				os.O_WRONLY|os.O_CREATE|os.O_TRUNC,
+				0600,
+			)
+
+			if err != nil {
+				return msg, err
+			}
+
+			l.channels[clientChannelID] = f
+
+			l.starttime = time.Now()
+
+			_, err = fmt.Fprintf(
+				f,
+				"{\"version\": 2, \"width\": %d, \"height\": %d, \"timestamp\": %d, \"env\": %v}\n",
+				l.initWidth,
+				l.initHeight,
+				l.starttime.Unix(),
+				string(jsonEnvs),
+			)
+
+			if err != nil {
+				return msg, err
+			}
+		}
+	}
+	return msg, nil
+}
+
+func (l *asciicastLogger) Close() (err error) {
+	for _, f := range l.channels {
+		_ = f.Close()
+	}
+	return nil
+}

cmd/sshpiperd/daemon.go

@@ -24,6 +24,8 @@ type daemon struct {
 	loginGraceTime time.Duration
 
 	recorddir             string
+	recordfmt             string
+	usernameAsRecorddir   bool
 	filterHostkeysReqeust bool
 }
 
@@ -218,23 +220,43 @@ func (d *daemon) run() error {
 			var downhook func([]byte) ([]byte, error)
 
 			if d.recorddir != "" {
-				recorddir := path.Join(d.recorddir, p.DownstreamConnMeta().User())
+				var recorddir string
+				if d.usernameAsRecorddir {
+					recorddir = path.Join(d.recorddir, p.DownstreamConnMeta().User())
+				} else {
+					uniqID := plugin.GetUniqueID(p.ChallengeContext())
+					recorddir = path.Join(d.recorddir, uniqID)
+				}
 				err = os.MkdirAll(recorddir, 0700)
 				if err != nil {
 					log.Errorf("cannot create screen recording dir %v: %v", recorddir, err)
 					return
 				}
+				if d.recordfmt == "asciicast" {
+					prefix := ""
+					if d.usernameAsRecorddir {
+						// add prefix to avoid conflict
+						prefix = fmt.Sprintf("%d-", time.Now().Unix())
+					}
+					recorder := newAsciicastLogger(recorddir, prefix)
+					defer recorder.Close()
 
-				recorder, err := newFilePtyLogger(recorddir)
-				if err != nil {
-					log.Errorf("cannot create screen recording logger: %v", err)
-					return
+					uphook = recorder.uphook
+					downhook = recorder.downhook
+				} else if d.recordfmt == "typescript" {
+					recorder, err := newFilePtyLogger(recorddir)
+					if err != nil {
+						log.Errorf("cannot create screen recording logger: %v", err)
+						return
+					}
+					defer recorder.Close()
+
+					uphook = recorder.loggingTty
 				}
-
-				uphook = recorder.loggingTty
 			}
 
 			if d.filterHostkeysReqeust {
+				nextUpHook := uphook
 				uphook = func(b []byte) ([]byte, error) {
 					if b[0] == 80 {
 						var x struct {
@@ -245,7 +267,7 @@ func (d *daemon) run() error {
 							return nil, nil
 						}
 					}
-					return b, nil
+					return nextUpHook(b)
 				}
 			}
 

cmd/sshpiperd/internal/plugin/grpc.go

@@ -592,3 +592,13 @@ func DialCmd(cmd *exec.Cmd) (*CmdPlugin, error) {
 
 	return &CmdPlugin{*g, ch}, nil
 }
+
+func GetUniqueID(ctx ssh.ChallengeContext) string {
+	switch meta := ctx.(type) {
+	case *connMeta:
+		return meta.UniqId
+	case *chainConnMeta:
+		return meta.UniqId
+	}
+	panic("unknown challenge context")
+}

cmd/sshpiperd/main.go

@@ -137,10 +137,22 @@ func main() {
 				EnvVars: []string{"SSHPIPERD_LOG_FORMAT"},
 			},
 			&cli.StringFlag{
-				Name:    "typescript-log-dir",
+				Name:    "screen-recording-dir",
 				Value:   "",
-				Usage:   "create typescript format screen recording and save into the directory see https://linux.die.net/man/1/script",
-				EnvVars: []string{"SSHPIPERD_TYPESCRIPT_LOG_DIR"},
+				Usage:   "the directory to save screen recording files",
+				EnvVars: []string{"SSHPIPERD_SCREEN_RECORDING_DIR"},
+			},
+			&cli.StringFlag{
+				Name:    "screen-recording-format",
+				Value:   "asciicast",
+				Usage:   "the format of screen recording files, one of: typescript (https://linux.die.net/man/1/script), asciicast (https://docs.asciinema.org/manual/asciicast/v2)",
+				EnvVars: []string{"SSHPIPERD_SCREEN_RECORDING_FORMAT"},
+			},
+			&cli.BoolFlag{
+				Name:    "username-as-recorddir",
+				Value:   false,
+				Usage:   "use the username as the directory name for saving screen recording files",
+				EnvVars: []string{"SSHPIPERD_USERNAME_AS_RECORDDIR"},
 			},
 			&cli.StringFlag{
 				Name:    "banner-text",
@@ -256,9 +268,15 @@ func main() {
 				return err
 			}
 
-			d.recorddir = ctx.String("typescript-log-dir")
+			d.recorddir = ctx.String("screen-recording-dir")
+			d.recordfmt = ctx.String("screen-recording-format")
+			d.usernameAsRecorddir = ctx.Bool("username-as-recorddir")
 			d.filterHostkeysReqeust = ctx.Bool("drop-hostkeys-message")
 
+			if d.recordfmt != "typescript" && d.recordfmt != "asciicast" {
+				return fmt.Errorf("invalid screen recording format: %v", d.recordfmt)
+			}
+
 			go func() {
 				quit <- d.run()
 			}()