handler.go

package k8ssshpod

import (
	"context"
	"fmt"
	"io"

	"github.com/gliderlabs/ssh"
	log "github.com/sirupsen/logrus"
	"golang.org/x/term"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
)

func ssHandler(client kubernetes.Interface, config *rest.Config, audit bool) func(s ssh.Session) {
	return func(s ssh.Session) {
		defer s.Exit(0)
		io.WriteString(s, "Wellcome to use K8S-SSH-POD\n")
		t := term.NewTerminal(s, "")

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		for {
			// select namespace
		select_namespace:
			lsns := func() ([]string, error) { return listNamespaces(ctx, client, nil) }
			ns, exit, err := promptAndSelectNames(s, t, lsns)
			if err != nil {
				log.Println(err)
				_, _ = fmt.Fprintf(s, "[ERROR]: %v", err)
				return
			}
			switch exit {
			case label_back:
				goto select_namespace
			case label_exit:
				return
			}

			// select pod
		select_pod:
			var getPod = func() ([]string, error) { return listPodWithNamespace(ctx, client, ns, nil) }
			podName, exit, err := promptAndSelectNames(s, t, getPod)
			if err != nil {
				log.Errorln(err)
				_, _ = fmt.Fprintf(s, "[ERROR]: %v", err)
				return
			}
			switch exit {
			case label_back:
				goto select_namespace
			case label_exit:
				return
			}

			// select container
			var getContainer = func() ([]string, error) { return listContainersForPod(ctx, client, ns, podName) }
			containerName, exit, err := promptAndSelectNames(s, t, getContainer)
			if err != nil {
				log.Errorln(err)
				_, _ = fmt.Fprintf(s, "[ERROR]: %v", err)
				return
			}
			switch exit {
			case label_back:
				goto select_pod
			case label_exit:
				return
			}
			_, window, isTty := s.Pty()

			if isTty {
				podReader, deferHandler := auditLog(s, audit)
				// defer deferHandler()
				// exec
				if err := exec(ctx, client, config, podReader, s, s, ns, podName, containerName, TerminalSizeQueue{window}, isTty, deferHandler); err != nil {
					log.Errorln(err)
					return
				}
			} else {
				io.WriteString(s, "No PTY requested.\n")
				s.Exit(1)
			}

			log.Debugln("exit ssh handler")
		}
	}
}