policy-generator/cmd/policygen/prometheus.go

package main

import (
	"fmt"
	v1 "k8s.io/api/core/v1"
	"k8s.io/klog/v2"
	"regexp"
	"slices"
	"strconv"
	"strings"
)

var standardPodLabels = map[string]func(pod *v1.Pod) string{
	"__meta_kubernetes_pod_name":            func(pod *v1.Pod) string { return pod.Name },
	"__meta_kubernetes_pod_ip":              func(pod *v1.Pod) string { return pod.Status.PodIP },
	"__meta_kubernetes_pod_node_name":       func(pod *v1.Pod) string { return pod.Spec.NodeName },
	"__meta_kubernetes_pod_uid":             func(pod *v1.Pod) string { return string(pod.UID) },
	"__meta_kubernetes_pod_phase":           func(pod *v1.Pod) string { return string(pod.Status.Phase) },
	"__meta_kubernetes_pod_ready":           func(pod *v1.Pod) string { return getPodReadyCondition(pod) },
	"__meta_kubernetes_pod_host_ip":         func(pod *v1.Pod) string { return pod.Status.HostIP },
	"__meta_kubernetes_namespace":           func(pod *v1.Pod) string { return pod.Namespace },
	"__meta_kubernetes_pod_container_init":  func(pod *v1.Pod) string { return "false" },
	"__meta_kubernetes_pod_service_account": func(pod *v1.Pod) string { return pod.Spec.ServiceAccountName },
	"__meta_kubernetes_pod_controller_kind": func(pod *v1.Pod) string { return getControllerKind(pod) },
	"__meta_kubernetes_pod_controller_name": func(pod *v1.Pod) string { return getControllerName(pod) },
}

var containerLabels = map[string]func(pod *v1.Pod, container *v1.Container) string{
	"__meta_kubernetes_pod_container_name": func(pod *v1.Pod, container *v1.Container) string {
		return container.Name
	},
	"__meta_kubernetes_pod_container_id": func(pod *v1.Pod, container *v1.Container) string {
		return getContainerID(pod, container)
	},
}

func metadataPodLabels(pod *v1.Pod) map[string]string {
	labels := make(map[string]string)
	for name, value := range pod.Labels {
		key := "__meta_kubernetes_pod_label_" + strings.ReplaceAll(name, ".", "_")
		key = strings.ReplaceAll(key, "/", "_")
		key = strings.ReplaceAll(key, "-", "_")
		labels[key] = value
	}
	for name, value := range pod.Annotations {
		key := "__meta_kubernetes_pod_annotation_" + strings.ReplaceAll(name, ".", "_")
		key = strings.ReplaceAll(key, "/", "_")
		key = strings.ReplaceAll(key, "-", "_")
		labels[key] = value
	}
	klog.V(6).Infof("METADATA LABELS %v\n", labels)
	return labels
}

func getContainerPortSourceLabels(port *v1.ContainerPort) map[string]string {
	labels := make(map[string]string)
	labels["__meta_kubernetes_pod_container_port_name"] = port.Name
	labels["__meta_kubernetes_pod_container_port_number"] = strconv.FormatInt(int64(port.ContainerPort), 10)
	labels["__meta_kubernetes_pod_container_port_protocol"] = string(port.Protocol)
	return labels
}

type RelabelingMatcher func(pod *v1.Pod) ([]v1.ContainerPort, error)

func MatchRelabeling(config RelabelConfig, pod *v1.Pod, container *v1.Container,
	port *v1.ContainerPort) (bool, error) {
	if config.Action != "keep" {
		return true, nil
	}
	labelValues := make([]string, len(config.SourceLabels))
	for isource, sourceLabel := range config.SourceLabels {
		if podLabelFunc, ok := standardPodLabels[sourceLabel]; ok {
			labelValues[isource] = podLabelFunc(pod)
			continue
		}
		metadataLabels := metadataPodLabels(pod)
		if metadataLabelValue, ok := metadataLabels[sourceLabel]; ok {
			labelValues[isource] = metadataLabelValue
			continue
		}
		if containerLabelFunc, ok := containerLabels[sourceLabel]; ok {
			labelValues[isource] = containerLabelFunc(pod, container)
			continue
		}
		containerPortSourceLabels := getContainerPortSourceLabels(port)
		if containerPortValue, ok := containerPortSourceLabels[sourceLabel]; ok {
			labelValues[isource] = containerPortValue
			continue
		}
	}
	// concatenate the label values
	labelstring := strings.Join(labelValues, ";")
	matched, err := regexp.MatchString(config.Regex, labelstring)
	if err != nil {
		return false, err
	}
	klog.V(5).Infof("      relabeling '%s' ~ '%s': %v\n", labelstring, config.Regex, matched)
	return matched, err
}

func combinations(s [][]string, i int, sequence []string) {
	if i == len(s) {
		fmt.Printf("COMBINATION %v\n", sequence)
		return
	}
	for _, v := range s[i] {
		sequence[i] = v
		combinations(s, i+1, sequence)
	}
}

// Helper functions
func getPodReadyCondition(pod *v1.Pod) string {
	for _, condition := range pod.Status.Conditions {
		if condition.Type == v1.PodReady {
			if condition.Status == v1.ConditionTrue {
				return "true"
			}
			return "false"
		}
	}
	return "unknown"
}

func getContainerID(pod *v1.Pod, container *v1.Container) string {
	containerIndex := slices.IndexFunc(pod.Spec.Containers, func(c v1.Container) bool {
		return c.Name == container.Name
	})
	id := pod.Status.ContainerStatuses[containerIndex].ContainerID
	// Strip the docker:// prefix if present
	return strings.TrimPrefix(id, "docker://")
}

func getControllerKind(pod *v1.Pod) string {
	for _, owner := range pod.OwnerReferences {
		if owner.Controller != nil && *owner.Controller {
			return owner.Kind
		}
	}
	return ""
}

func getControllerName(pod *v1.Pod) string {
	for _, owner := range pod.OwnerReferences {
		if owner.Controller != nil && *owner.Controller {
			return owner.Name
		}
	}
	return ""
}