k8s.go

package k8ssshpod

import (
	"context"
	"io"

	log "github.com/sirupsen/logrus"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
	"k8s.io/client-go/tools/remotecommand"
	"k8s.io/kubectl/pkg/scheme"
)

func NewK8SClient(inK8sCluster bool, kubeConfigFile string) (client *kubernetes.Clientset, config *rest.Config, err error) {
	if inK8sCluster {
		config, err = rest.InClusterConfig()
	} else {
		kubeConfigFile = RealPath(kubeConfigFile)
		log.Infof("Initialize k8s client with config file: %s", kubeConfigFile)
		config, err = clientcmd.BuildConfigFromFlags("", kubeConfigFile)
	}
	if err != nil {
		return
	}
	clientset, err := kubernetes.NewForConfig(config)
	return clientset, config, err
}

func namespaces(ctx context.Context, client kubernetes.Interface, opt metav1.ListOptions) (*corev1.NamespaceList, error) {
	return client.CoreV1().Namespaces().List(ctx, opt)
}

func pods(ctx context.Context, client kubernetes.Interface, namespace string, opt metav1.ListOptions) (*corev1.PodList, error) {
	return client.CoreV1().Pods(namespace).List(ctx, opt)
}

func exec(ctx context.Context,
	k8sClient kubernetes.Interface,
	cfg *rest.Config,
	stdin io.Reader,
	stdout io.Writer,
	stderr io.Writer,
	namespace,
	podName,
	containerName string,
	terminalSizeQueue remotecommand.TerminalSizeQueue,
	isTty bool,
	deferHandler func(),
) error {
	if deferHandler != nil {
		defer deferHandler()
	}

	if len(containerName) < 1 {
		pod, err := k8sClient.CoreV1().Pods(namespace).Get(ctx, podName, metav1.GetOptions{})
		if err != nil {
			log.Errorln(err)
			return err
		}
		containerName = pod.Spec.Containers[0].Name
	}
	req := k8sClient.CoreV1().RESTClient().Post().
		Resource("pods").
		Name(podName).
		Namespace(namespace).
		SubResource("exec")

	req.VersionedParams(&corev1.PodExecOptions{
		Container: containerName,
		Command:   []string{"/bin/sh"},
		Stdin:     true,
		Stdout:    true,
		Stderr:    true,
		TTY:       isTty,
	}, scheme.ParameterCodec)

	exec, err := remotecommand.NewSPDYExecutor(cfg, "POST", req.URL())
	if err != nil {
		log.Error(err)
		return err
	}

	opt := remotecommand.StreamOptions{
		Stdin:             stdin,
		Stdout:            stdout,
		Stderr:            stderr,
		TerminalSizeQueue: terminalSizeQueue,
		Tty:               isTty,
	}
	if err := exec.StreamWithContext(ctx, opt); err != nil {
		log.Error(err)
		return err
	}

	return nil
}