node.go

package node

import (
	"context"
	"encoding/json"
	"sync"

	"github.com/aquasecurity/trivy/pkg/report"
	"github.com/prometheus/client_golang/prometheus"
	"golang.org/x/sync/semaphore"
	"golang.org/x/xerrors"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/klog/v2"
)

var (
	nodeVulnerbilities = prometheus.NewGaugeVec(prometheus.GaugeOpts{
		Namespace: "trivy",
		Name:      "node_vulnerabilities",
		Help:      "Vulnerbilities of cluster node detected by trivy",
	}, []string{
		"nodeName",
		"vulnerabilityId",
		"pkgName",
		"installedVersion",
		"severity",
		"fixedVersion",
	},
	)
)

type NodeScanner struct {
	kubernetesClient KubernetesClient
	namespace        string
	concurrency      uint8
}

func NewNodeScanner(kclient KubernetesClient, ns string, concurrency uint8) *NodeScanner {
	return &NodeScanner{
		kubernetesClient: kclient,
		namespace:        ns,
		concurrency:      concurrency,
	}
}

func (n *NodeScanner) Name() string {
	return "node-vulnerability-scanner"
}

func (n *NodeScanner) Scan(ctx context.Context) (prometheus.Collector, error) {
	nodes, err := n.kubernetesClient.ListNodes(ctx)
	if err != nil {
		return nodeVulnerbilities, xerrors.Errorf("failed to list nodes for scanning: %w", err)
	}

	wg := sync.WaitGroup{}
	reports := map[string]report.Report{}
	sem := semaphore.NewWeighted(int64(n.concurrency))
	for _, node := range nodes {
		wg.Add(1)
		if err := sem.Acquire(ctx, 1); err != nil {
			klog.Warningf("failed to acquire semaphore: %w", err)
			continue
		}

		go func(nodeName string) {
			defer func() {
				sem.Release(1)
				wg.Done()
			}()
			report, err := n.ScanNode(ctx, nodeName)
			if err != nil {
				klog.Warningf("failed to scan node(%s): %s", nodeName, err)
				return
			}

			reports[nodeName] = *report
		}(node.Name)
	}
	wg.Wait()

	nodeVulnerbilities.Reset()
	for nodeName, rp := range reports {
		for _, rs := range rp.Results {
			for _, v := range rs.Vulnerabilities {
				labels := []string{
					nodeName,
					v.VulnerabilityID,
					v.PkgName,
					v.InstalledVersion,
					v.Severity,
					v.FixedVersion,
				}
				nodeVulnerbilities.WithLabelValues(labels...).Set(1)
			}
		}
	}

	return nodeVulnerbilities, nil
}

func (n *NodeScanner) CleanUp() error {
	labelSelector := &metav1.LabelSelector{
		MatchLabels: map[string]string{
			"app": "trivy-node-scanner",
		},
	}

	pods, err := n.kubernetesClient.ListPodsByLabel(context.Background(), labelSelector)
	if err != nil {
		return xerrors.Errorf("failed to list the remaining pods scanning nodes: %s", err)
	}

	for _, p := range pods {
		err := n.kubernetesClient.DeletePod(context.Background(), p.Namespace, p.Name)
		if err != nil {
			return xerrors.Errorf("failed to delete pod(%s): %s", p.Name, err)
		}
	}

	return nil
}

func (n *NodeScanner) ScanNode(ctx context.Context, nodeName string) (*report.Report, error) {
	podName := "trivy-" + nodeName
	pod := buildPod(podName, nodeName)

	err := n.kubernetesClient.CreatePod(ctx, n.namespace, pod)
	if err != nil {
		return nil, xerrors.Errorf("failed to create pod(%s) for scanning node: %s", pod.Name, err)
	}

	defer func() {
		err = n.kubernetesClient.DeletePod(context.Background(), n.namespace, pod.Name)
		if err != nil {
			klog.Warningf("failed to delete pod(%s) after scanning %s: %s", pod.Name, nodeName, err)
		}
	}()

	err = n.kubernetesClient.WaitUntilCompleted(n.namespace, pod.Name)
	if err != nil {
		return nil, xerrors.Errorf("failed to complete pod(%s) for scanning node: %s", pod.Name, err)
	}

	containerName := "trivy"
	log, err := n.kubernetesClient.GetContainerLog(ctx, n.namespace, pod.Name, containerName)
	if err != nil {
		return nil, xerrors.Errorf("failed to get container log from %s: %w", pod.Name, err)
	}

	var rp report.Report
	err = json.Unmarshal([]byte(log), &rp)
	if err != nil {
		return nil, xerrors.Errorf("failed to unmarshal from container log string to report: %w", err)
	}

	return &rp, nil
}

func buildPod(podName string, nodeName string) *v1.Pod {
	return &v1.Pod{
		TypeMeta: metav1.TypeMeta{
			Kind:       "Pod",
			APIVersion: "v1",
		},
		ObjectMeta: metav1.ObjectMeta{
			Name: podName,
			Labels: map[string]string{
				"app": "trivy-node-scanner",
			},
		},
		Spec: v1.PodSpec{
			Affinity: &v1.Affinity{
				NodeAffinity: &v1.NodeAffinity{
					RequiredDuringSchedulingIgnoredDuringExecution: &v1.NodeSelector{
						NodeSelectorTerms: []v1.NodeSelectorTerm{
							{
								MatchExpressions: []v1.NodeSelectorRequirement{
									{
										Key:      "kubernetes.io/hostname",
										Operator: v1.NodeSelectorOpIn,
										Values:   []string{nodeName},
									},
								},
							},
						},
					},
				},
			},
			Tolerations: []v1.Toleration{
				{
					Key:      "node-role.kubernetes.io/master",
					Effect:   v1.TaintEffectNoSchedule,
					Operator: v1.TolerationOpExists,
				},
				{
					Key:      "node-role.kubernetes.io/control-plane",
					Effect:   v1.TaintEffectNoSchedule,
					Operator: v1.TolerationOpExists,
				},
			},
			Containers: []v1.Container{
				{
					Name:  "trivy",
					Image: "ghcr.io/aquasecurity/trivy:0.22.0",
					Args: []string{
						"--quiet",
						"rootfs",
						"--format",
						"json",
						"--no-progress",
						"--skip-dirs",
						"/host/proc",
						"/host",
					},
					VolumeMounts: []v1.VolumeMount{
						{
							Name:      "host",
							MountPath: "/host",
						},
					},
				},
			},
			RestartPolicy: v1.RestartPolicyNever,
			Volumes: []v1.Volume{
				{
					Name: "host",
					VolumeSource: v1.VolumeSource{
						HostPath: &v1.HostPathVolumeSource{
							Path: "/",
						},
					},
				},
			},
		},
	}
}