From bf944cb048b1b505c1ace7f2073254d449da9f69 Mon Sep 17 00:00:00 2001
From: nhyatt <nhyatt@smoothnet.org>
Date: Fri, 24 Mar 2023 13:06:58 -0500
Subject: [PATCH] properly mutate the container

---
 internal/operations/podsMutation.go | 58 +++++++++++++++++++++++++----
 1 file changed, 50 insertions(+), 8 deletions(-)

diff --git a/internal/operations/podsMutation.go b/internal/operations/podsMutation.go
index 238a7c0..d82ef3b 100644
--- a/internal/operations/podsMutation.go
+++ b/internal/operations/podsMutation.go
@@ -2,6 +2,8 @@ package operations
 
 import (
 	"fmt"
+	"log"
+	"regexp"
 
 	admission "k8s.io/api/admission/v1"
 	core "k8s.io/api/core/v1"
@@ -27,21 +29,46 @@ func PodsMutation() Hook {
 
 func podMutationCreate() AdmitFunc {
 	return func(r *admission.AdmissionRequest, cfg *config.Config) (*Result, error) {
-		var operations []PatchOperation
+		var (
+			operations []PatchOperation
+			mutated    bool
+		)
+
 		pod, err := parsePod(r.Object.Raw)
 		if err != nil {
 			return &Result{Msg: err.Error()}, nil
 		}
 
 		// if pod is administratively exempt
-		if cfg.AllowAdminNoMutate && func(pod *core.Pod) bool {
-			for label, value := range pod.Annotations {
-				if label == "AdminNoMutate" && value == "true" {
-					return false
+		if func(serviceEnabled bool, pod *core.Pod) bool {
+			if serviceEnabled {
+				for label, value := range pod.Annotations {
+					if label == "AdminNoMutate" && value == "true" {
+						return false
+					}
 				}
 			}
 			return true
-		}(pod) {
+		}(cfg.AllowAdminNoMutate, pod) {
+			for i, p := range pod.Spec.Containers {
+				img, mutationOccurred, err := customDockerRegistry(p.Image, cfg)
+				if err != nil {
+					return &Result{Msg: err.Error()}, nil
+				}
+				if mutationOccurred {
+					mutated = true
+					path := fmt.Sprintf("/spec/containers/%d/image", i)
+					operations = append(operations, ReplacePatchOperation(path, img))
+					log.Printf("[TRACE] Image has been mutated: %s -> %s", p.Image, img)
+				} else {
+					log.Printf("[TRACE] No mutation required for image: %s", p.Image)
+				}
+			}
+		} else {
+			log.Printf("[TRACE] Mutations administratively disabled.")
+		}
+
+		if mutated {
 			// mutate pod (annotation)
 			metadata := map[string]string{
 				"mutation-status": "pod mutated by mutation-controller",
@@ -52,8 +79,6 @@ func podMutationCreate() AdmitFunc {
 			}
 			// add annotation stating that the pos had been mutated
 			operations = append(operations, AddPatchOperation("/metadata/annotations", metadata))
-
-			// add image mutation
 		}
 
 		return &Result{
@@ -62,3 +87,20 @@ func podMutationCreate() AdmitFunc {
 		}, nil
 	}
 }
+
+func customDockerRegistry(imgPath string, cfg *config.Config) (string, bool, error) {
+	if len(cfg.DockerhubRegistry) == 0 {
+		return imgPath, false, nil
+	}
+
+	// regex match official project
+	reg, err := regexp.Compile(`^([a-z]|\.|_|-)+\:([a-zA-Z0-9]|_|\.|-)+$`)
+	if err != nil {
+		return "", false, fmt.Errorf("Unable to parse regex: %v", err)
+	}
+	if reg.MatchString(imgPath) {
+		log.Printf("Official docker image detected: %s", imgPath)
+		return fmt.Sprintf("%s/library/%s", cfg.DockerhubRegistry, imgPath), true, nil
+	}
+	return "", false, nil
+}