oc-monitord/tools/kubernetes.go

package tools

import (
	"context"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"oc-monitord/conf"
	"oc-monitord/models"
	"oc-monitord/utils"
	"os"
	"sync"
	"time"

	"cloud.o-forge.io/core/oc-lib/models/common/enum"
	wfv1 "github.com/argoproj/argo-workflows/v3/pkg/apis/workflow/v1alpha1"
	"github.com/argoproj/argo-workflows/v3/pkg/client/clientset/versioned"
	"github.com/google/uuid"
	corev1 "k8s.io/api/core/v1"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/serializer"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"
)

type KubernetesTools struct {
	Set          *kubernetes.Clientset
	VersionedSet *versioned.Clientset
}

func NewKubernetesTool() (Tool, error) {
	// Load Kubernetes config (from ~/.kube/config)
	config := &rest.Config{
		Host: conf.GetConfig().KubeHost + ":" + conf.GetConfig().KubePort,
		TLSClientConfig: rest.TLSClientConfig{
			CAData:   []byte(conf.GetConfig().KubeCA),
			CertData: []byte(conf.GetConfig().KubeCert),
			KeyData:  []byte(conf.GetConfig().KubeData),
		},
	}

	// Create clientset
	clientset, err := kubernetes.NewForConfig(config)
	if err != nil {
		return nil, errors.New("Error creating Kubernetes client: " + err.Error())
	}

	clientset2, err := versioned.NewForConfig(config)
	if err != nil {
		return nil, errors.New("Error creating Kubernetes versionned client: " + err.Error())
	}
	return &KubernetesTools{
		Set:          clientset,
		VersionedSet: clientset2,
	}, nil
}

func (k *KubernetesTools) LogWorkflow(execID string, namespace string, workflowName string, argoFilePath string, stepMax int, current_watch *models.ArgoWatch, previous_watch *models.ArgoWatch, argoLogs *models.ArgoLogs,
	seen []string, logFunc func(argoFilePath string, stepMax int, pipe io.ReadCloser, current_watch *models.ArgoWatch, previous_watch *models.ArgoWatch, argoLogs *models.ArgoLogs, seen []string, wg *sync.WaitGroup)) error {
	exec := utils.GetExecution(execID)
	if exec == nil {
		return errors.New("Could not retrieve workflow ID from execution ID " + execID)
	}
	if exec.State == enum.DRAFT || exec.State == enum.FAILURE || exec.State == enum.SUCCESS {
		return nil
	}
	k.logWorkflow(namespace, workflowName, argoFilePath, stepMax, current_watch, previous_watch, argoLogs, seen, logFunc)
	return k.LogWorkflow(execID, namespace, workflowName, argoFilePath, stepMax, current_watch, previous_watch, argoLogs, seen, logFunc)
}

func (k *KubernetesTools) logWorkflow(namespace string, workflowName string, argoFilePath string, stepMax int, current_watch *models.ArgoWatch, previous_watch *models.ArgoWatch, argoLogs *models.ArgoLogs,
	seen []string,
	logFunc func(argoFilePath string, stepMax int, pipe io.ReadCloser, current_watch *models.ArgoWatch, previous_watch *models.ArgoWatch, argoLogs *models.ArgoLogs, seen []string, wg *sync.WaitGroup)) error {
	// List pods related to the Argo workflow
	labelSelector := fmt.Sprintf("workflows.argoproj.io/workflow=%s", workflowName)
	for retries := 0; retries < 10; retries++ { // Retry for up to ~20 seconds
		// List workflow pods
		wfPods, err := k.Set.CoreV1().Pods(namespace).List(context.Background(), metav1.ListOptions{
			LabelSelector: labelSelector,
		})
		if err != nil {
			return err
		}
		// If we found pods, stream logs
		if len(wfPods.Items) > 0 {
			var wg sync.WaitGroup
			// Stream logs from all matching pods
			for _, pod := range wfPods.Items {
				for _, container := range pod.Spec.Containers {
					wg.Add(1)
					go k.streamLogs(namespace, pod.Name, container.Name, argoFilePath, stepMax, &wg, current_watch, previous_watch, argoLogs, seen, logFunc)
				}
			}
			wg.Wait()
			return nil
		}
		time.Sleep(2 * time.Second) // Wait before retrying
	}
	return errors.New("no pods found for the workflow")
}

// Function to stream logs
func (k *KubernetesTools) streamLogs(namespace string, podName string, containerName string,
	argoFilePath string, stepMax int, wg *sync.WaitGroup, current_watch *models.ArgoWatch, previous_watch *models.ArgoWatch, argoLogs *models.ArgoLogs, seen []string,
	logFunc func(argo_file_path string, stepMax int, pipe io.ReadCloser, current_watch *models.ArgoWatch, previous_watch *models.ArgoWatch, argoLogs *models.ArgoLogs, seen []string, wg *sync.WaitGroup)) {
	req := k.Set.CoreV1().Pods(namespace).GetLogs(podName, &corev1.PodLogOptions{
		Container: containerName, // Main container
		Follow:    true,          // Equivalent to -f flag in kubectl logs
	})
	defer wg.Done()
	// Open stream
	stream, err := req.Stream(context.Background())
	if err != nil {
		return
	}
	defer stream.Close()
	var internalWg sync.WaitGroup
	logFunc(argoFilePath, stepMax, stream, current_watch, previous_watch, argoLogs, seen, &internalWg)
	internalWg.Wait()
}

func (k *KubernetesTools) CreateArgoWorkflow(path string, ns string) (string, error) {
	// Read workflow YAML file
	workflowYAML, err := os.ReadFile(path)
	if err != nil {
		return "", err
	}
	// Decode the YAML into a Workflow struct
	scheme := runtime.NewScheme()
	_ = wfv1.AddToScheme(scheme)
	codecs := serializer.NewCodecFactory(scheme)
	decode := codecs.UniversalDeserializer().Decode

	obj, _, err := decode(workflowYAML, nil, nil)
	if err != nil {
		return "", errors.New("failed to decode YAML: " + err.Error())
	}

	workflow, ok := obj.(*wfv1.Workflow)
	if !ok {
		return "", errors.New("decoded object is not a Workflow")
	}

	// Create the workflow in the "argo" namespace
	createdWf, err := k.VersionedSet.ArgoprojV1alpha1().Workflows(ns).Create(context.Background(), workflow, metav1.CreateOptions{})
	if err != nil {
		return "", errors.New("failed to create workflow: " + err.Error())
	}
	fmt.Printf("workflow %s created in namespace %s\n", createdWf.Name, "argo")
	return createdWf.Name, nil
}

func (k *KubernetesTools) CreateAccessSecret(ns string, login string, password string) (string, error) {
	// Namespace where the secret will be created
	namespace := "default"
	// Encode the secret data (Kubernetes requires base64-encoded values)
	secretData := map[string][]byte{
		"access-key": []byte(base64.StdEncoding.EncodeToString([]byte(login))),
		"secret-key": []byte(base64.StdEncoding.EncodeToString([]byte(password))),
	}

	// Define the Secret object
	name := uuid.New().String()
	secret := &v1.Secret{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: ns,
		},
		Type: v1.SecretTypeOpaque,
		Data: secretData,
	}
	// Create the Secret in Kubernetes
	_, err := k.Set.CoreV1().Secrets(namespace).Create(context.Background(), secret, metav1.CreateOptions{})
	if err != nil {
		return "", errors.New("Error creating secret: " + err.Error())
	}
	return name, nil
}