Extract images from PodSpec using Cue

cue.mod/module.cue

@@ -0,0 +1,4 @@
+module: "github.com/mpv/kir"
+language: {
+	version: "v0.9.0"
+}

cueparser/cueparser.go

@@ -0,0 +1,231 @@
+package cueparser
+
+import (
+	"fmt"
+	"strings"
+
+	"cuelang.org/go/cue"
+	"cuelang.org/go/cue/cuecontext"
+	"cuelang.org/go/cue/load"
+	"cuelang.org/go/encoding/yaml"
+)
+
+// ProcessData processes YAML data and extracts container images from PodSpec
+func ProcessData(data []byte) ([]string, error) {
+	// Create a CUE context
+	ctx := cuecontext.New()
+
+	// Load the PodSpec schema from the CUE Central Registry
+	bis := load.Instances([]string{"cue.dev/x/k8s.io/api/core/v1"}, nil)
+	if len(bis) == 0 {
+		// Fall back to local schema if Central Registry is not available
+		return processWithLocalSchema(data)
+	}
+
+	pkgV := ctx.BuildInstance(bis[0])
+	if pkgV.Err() != nil {
+		// Fall back to local schema if Central Registry is not available
+		return processWithLocalSchema(data)
+	}
+
+	podSpec := pkgV.LookupPath(cue.ParsePath("#PodSpec"))
+	if podSpec.Err() != nil {
+		// Fall back to local schema if Central Registry is not available
+		return processWithLocalSchema(data)
+	}
+
+	// Load the YAML data
+	dataV, err := yaml.Extract("", data)
+	if err != nil {
+		return nil, fmt.Errorf("failed to extract YAML: %v", err)
+	}
+
+	dataValue := ctx.BuildFile(dataV)
+	if dataValue.Err() != nil {
+		return nil, fmt.Errorf("failed to build YAML data: %v", dataValue.Err())
+	}
+
+	// Unify the YAML value with the schema and validate
+	combined := podSpec.Unify(dataValue)
+	if err := combined.Validate(cue.Concrete(true)); err != nil {
+		return nil, fmt.Errorf("validation error: %v", err)
+	}
+
+	// Extract container images
+	var images []string
+
+	// Try to get containers from the validated PodSpec
+	containersValue := combined.LookupPath(cue.ParsePath("containers"))
+	if containersValue.Exists() {
+		// Iterate through containers
+		iter, err := containersValue.List()
+		if err != nil {
+			return nil, fmt.Errorf("failed to iterate containers: %v", err)
+		}
+
+		for iter.Next() {
+			container := iter.Value()
+			imageValue := container.LookupPath(cue.ParsePath("image"))
+			if imageValue.Exists() {
+				image, err := imageValue.String()
+				if err != nil {
+					return nil, fmt.Errorf("failed to get image string: %v", err)
+				}
+				images = append(images, image)
+			}
+		}
+	}
+
+	// Try to get initContainers from the validated PodSpec
+	initContainersValue := combined.LookupPath(cue.ParsePath("initContainers"))
+	if initContainersValue.Exists() {
+		// Iterate through initContainers
+		iter, err := initContainersValue.List()
+		if err != nil {
+			return nil, fmt.Errorf("failed to iterate initContainers: %v", err)
+		}
+
+		for iter.Next() {
+			container := iter.Value()
+			imageValue := container.LookupPath(cue.ParsePath("image"))
+			if imageValue.Exists() {
+				image, err := imageValue.String()
+				if err != nil {
+					return nil, fmt.Errorf("failed to get image string: %v", err)
+				}
+				images = append(images, image)
+			}
+		}
+	}
+
+	return images, nil
+}
+
+// processWithLocalSchema processes YAML data using a local schema definition
+func processWithLocalSchema(data []byte) ([]string, error) {
+	// Create a CUE context
+	ctx := cuecontext.New()
+
+	// Load the local PodSpec schema
+	schema := `
+	#PodSpec: {
+		containers?: [...#Container]
+		initContainers?: [...#Container]
+	}
+
+	#Container: {
+		name: string
+		image: string
+	}
+	`
+	schemaValue := ctx.CompileString(schema)
+	if schemaValue.Err() != nil {
+		return nil, fmt.Errorf("failed to compile schema: %v", schemaValue.Err())
+	}
+
+	// Load the YAML data
+	dataV, err := yaml.Extract("", data)
+	if err != nil {
+		return nil, fmt.Errorf("failed to extract YAML: %v", err)
+	}
+
+	dataValue := ctx.BuildFile(dataV)
+	if dataValue.Err() != nil {
+		return nil, fmt.Errorf("failed to build YAML data: %v", dataValue.Err())
+	}
+
+	// Unify the YAML value with the schema and validate
+	combined := schemaValue.LookupPath(cue.ParsePath("#PodSpec")).Unify(dataValue)
+	if err := combined.Validate(cue.Concrete(true)); err != nil {
+		return nil, fmt.Errorf("validation error: %v", err)
+	}
+
+	// Extract container images
+	var images []string
+
+	// Try to get containers from the validated PodSpec
+	containersValue := combined.LookupPath(cue.ParsePath("containers"))
+	if containersValue.Exists() {
+		// Iterate through containers
+		iter, err := containersValue.List()
+		if err != nil {
+			return nil, fmt.Errorf("failed to iterate containers: %v", err)
+		}
+
+		for iter.Next() {
+			container := iter.Value()
+			imageValue := container.LookupPath(cue.ParsePath("image"))
+			if imageValue.Exists() {
+				image, err := imageValue.String()
+				if err != nil {
+					return nil, fmt.Errorf("failed to get image string: %v", err)
+				}
+				images = append(images, image)
+			}
+		}
+	}
+
+	// Try to get initContainers from the validated PodSpec
+	initContainersValue := combined.LookupPath(cue.ParsePath("initContainers"))
+	if initContainersValue.Exists() {
+		// Iterate through initContainers
+		iter, err := initContainersValue.List()
+		if err != nil {
+			return nil, fmt.Errorf("failed to iterate initContainers: %v", err)
+		}
+
+		for iter.Next() {
+			container := iter.Value()
+			imageValue := container.LookupPath(cue.ParsePath("image"))
+			if imageValue.Exists() {
+				image, err := imageValue.String()
+				if err != nil {
+					return nil, fmt.Errorf("failed to get image string: %v", err)
+				}
+				images = append(images, image)
+			}
+		}
+	}
+
+	return images, nil
+}
+
+// ProcessKubernetesYAML processes a Kubernetes YAML document and extracts container images
+func ProcessKubernetesYAML(data []byte) ([]string, error) {
+	// First, try to extract the PodSpec directly
+	images, err := ProcessData(data)
+	if err == nil && len(images) > 0 {
+		return images, nil
+	}
+
+	// If that fails, try to extract the PodSpec from a Kubernetes resource
+	// This is a simplified approach - in a real implementation, you would need to
+	// handle different Kubernetes resource types (Deployment, StatefulSet, etc.)
+
+	// For now, we'll just return the error from the first attempt
+	return nil, fmt.Errorf("failed to extract PodSpec: %v", err)
+}
+
+// ProcessKubernetesListYAML processes a Kubernetes List YAML document and extracts container images
+func ProcessKubernetesListYAML(data []byte) ([]string, error) {
+	// Split the YAML document by "---" to handle multiple resources
+	docs := strings.Split(string(data), "---")
+
+	var allImages []string
+	for _, doc := range docs {
+		if strings.TrimSpace(doc) == "" {
+			continue
+		}
+
+		images, err := ProcessKubernetesYAML([]byte(doc))
+		if err != nil {
+			// Log the error but continue processing other documents
+			fmt.Printf("Error processing document: %v\n", err)
+			continue
+		}
+
+		allImages = append(allImages, images...)
+	}
+
+	return allImages, nil
+}

cueparser/cueparser_test.go

@@ -0,0 +1,138 @@
+package cueparser
+
+import (
+	"testing"
+)
+
+func TestProcessData(t *testing.T) {
+	tests := []struct {
+		name    string
+		data    []byte
+		want    []string
+		wantErr bool
+	}{
+		{
+			name: "Valid PodSpec with containers",
+			data: []byte(`
+containers:
+  - name: test-container
+    image: test-image
+`),
+			want:    []string{"test-image"},
+			wantErr: false,
+		},
+		{
+			name: "Valid PodSpec with initContainers",
+			data: []byte(`
+initContainers:
+  - name: init-container
+    image: init-image
+`),
+			want:    []string{"init-image"},
+			wantErr: false,
+		},
+		{
+			name: "Valid PodSpec with both containers and initContainers",
+			data: []byte(`
+containers:
+  - name: test-container
+    image: test-image
+initContainers:
+  - name: init-container
+    image: init-image
+`),
+			want:    []string{"test-image", "init-image"},
+			wantErr: false,
+		},
+		{
+			name: "Invalid PodSpec (missing required fields)",
+			data: []byte(`
+containers:
+  - image: test-image
+`),
+			want:    nil,
+			wantErr: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := ProcessData(tt.data)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("ProcessData() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if !tt.wantErr {
+				if len(got) != len(tt.want) {
+					t.Errorf("ProcessData() got %v images, want %v", len(got), len(tt.want))
+					return
+				}
+				for i, img := range got {
+					if img != tt.want[i] {
+						t.Errorf("ProcessData() got image %v, want %v", img, tt.want[i])
+					}
+				}
+			}
+		})
+	}
+}
+
+func TestProcessKubernetesListYAML(t *testing.T) {
+	tests := []struct {
+		name    string
+		data    []byte
+		want    []string
+		wantErr bool
+	}{
+		{
+			name: "Multiple PodSpecs",
+			data: []byte(`
+---
+containers:
+  - name: container1
+    image: image1
+---
+containers:
+  - name: container2
+    image: image2
+`),
+			want:    []string{"image1", "image2"},
+			wantErr: false,
+		},
+		{
+			name: "Mixed valid and invalid PodSpecs",
+			data: []byte(`
+---
+containers:
+  - name: container1
+    image: image1
+---
+containers:
+  - image: image2
+`),
+			want:    []string{"image1"},
+			wantErr: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := ProcessKubernetesListYAML(tt.data)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("ProcessKubernetesListYAML() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if !tt.wantErr {
+				if len(got) != len(tt.want) {
+					t.Errorf("ProcessKubernetesListYAML() got %v images, want %v", len(got), len(tt.want))
+					return
+				}
+				for i, img := range got {
+					if img != tt.want[i] {
+						t.Errorf("ProcessKubernetesListYAML() got image %v, want %v", img, tt.want[i])
+					}
+				}
+			}
+		})
+	}
+}

cueparser/schema.cue

@@ -0,0 +1,14 @@
+package cueparser
+
+// PodSpec is a simplified schema for Kubernetes PodSpec
+#PodSpec: {
+	containers?: [...#Container]
+	initContainers?: [...#Container]
+}
+
+// Container represents a container in a pod
+#Container: {
+	name: string
+	image: string
+	// Add other fields as needed
+}