prompt_select.go

package k8ssshpod

import (
	"context"
	"fmt"
	"io"
	"strconv"

	"github.com/gliderlabs/ssh"
	log "github.com/sirupsen/logrus"
	"golang.org/x/term"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
)

type ListNamespaceHandler func(context.Context, kubernetes.Interface, metav1.ListOptions) (*corev1.NamespaceList, error)
type ListPodHandler func(context.Context, kubernetes.Interface, string, metav1.ListOptions) (*corev1.PodList, error)

const (
	label_back uint8 = 1
	label_exit uint8 = 2
)

func promptAndSelectNames(sess ssh.Session, term *term.Terminal, getNames func() ([]string, error)) (string, uint8, error) {
	names, err := getNames()
	if err != nil {
		return "", 0, err
	}

	// prompt
	for k, v := range names {
		_, _ = fmt.Fprintf(sess, "[%d] %s \r\n", k+1, v)
	}
	_, _ = fmt.Fprintln(sess, "Please select ID: ")
	sel, err := term.ReadLine()
	if err != nil {
		return "", 0, err
	}

	switch sel {
	case "quit", "exit":
		return "", label_exit, nil
	case "", "back", "b":
		return "", label_back, nil
	}

	i, err := strconv.Atoi(sel)
	if err != nil {
		io.WriteString(sess, fmt.Sprintf("Invalid input [%s], please again \r\n", sel))
		return promptAndSelectNames(sess, term, getNames)
	}

	if (i < 1) || (i > len(names)) {
		fmt.Fprintf(sess, "ID: [%d] out of range, please again\r\n", i)
		return promptAndSelectNames(sess, term, getNames)
	}

	return names[i-1], 0, nil
}

func listNamespaces(ctx context.Context, client kubernetes.Interface, listNamespaceHandler ListNamespaceHandler) ([]string, error) {
	if listNamespaceHandler == nil {
		listNamespaceHandler = namespaces
	}
	nsList, err := listNamespaceHandler(ctx, client, metav1.ListOptions{})
	if err != nil {
		return nil, err
	}
	nsNames := []string{}
	for _, item := range nsList.Items {
		nsNames = append(nsNames, item.Name)
	}

	return nsNames, nil
}

func listPodWithNamespace(ctx context.Context, client kubernetes.Interface, ns string, listPodHandler ListPodHandler) ([]string, error) {
	if listPodHandler != nil {
		listPodHandler = pods
	}
	podList, err := pods(ctx, client, ns, metav1.ListOptions{})
	if err != nil {
		return nil, err
	}

	podNames := []string{}
	for _, v := range podList.Items {
		podNames = append(podNames, v.Name)
	}

	return podNames, nil
}

func listContainersForPod(ctx context.Context, client kubernetes.Interface, ns string, podName string) ([]string, error) {
	pod, err := client.CoreV1().Pods(ns).Get(ctx, podName, metav1.GetOptions{})
	if err != nil {
		log.Errorln(err)
		return nil, err
	}

	names := []string{}
	for _, c := range pod.Spec.Containers {
		names = append(names, c.Name)
	}
	return names, nil
}