Merge pull request #18 from slashben/main

Adding finalization for objects
kubescape · Nov 13, 2023 · 5c37347 · 5c37347
2 parents 742ec1d + 3be895d
commit 5c37347
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 3 deletions.
diff --git a/main.go b/main.go
@@ -102,6 +102,7 @@ func main() {
 		EventSink:      eventSink,
 		Tracer:         tracer,
 		Interval:       60, // 60 seconds for now, TODO: make it configurable
+		FinalizeTime:   0,  // 0 seconds to disable finalization
 		K8sConfig:      k8sConfig,
 		RecordStrategy: collector.RecordStrategyOnlyIfNotExists,
 	}

diff --git a/pkg/collector/collector.go b/pkg/collector/collector.go
@@ -62,6 +62,8 @@ type CollectorManagerConfig struct {
 	EventSink *eventsink.EventSink
 	// Interval in seconds for collecting data from containers
 	Interval uint64
+	// Finalize application profiles time
+	FinalizeTime uint64
 	// Kubernetes configuration
 	K8sConfig *rest.Config
 	// Tracer object
@@ -106,7 +108,7 @@ func (cm *CollectorManager) ContainerStarted(id *ContainerId) {
 	// Check if applicaton profile already exists
 	appProfileExists, err := cm.doesApplicationProfileExists(id.Namespace, id.PodName, true, true)
 	if err != nil {
-		log.Printf("error checking if application profile exists: %s\n", err)
+		//log.Printf("error checking if application profile exists: %s\n", err)
 	} else if appProfileExists {
 		// If application profile exists, check if record strategy is RecordStrategyOnlyIfNotExists
 		if cm.config.RecordStrategy == RecordStrategyOnlyIfNotExists {
@@ -128,6 +130,11 @@ func (cm *CollectorManager) ContainerStarted(id *ContainerId) {
 
 	// Add a timer for collection of data from container events
 	startContainerTimer(id, cm.config.Interval, cm.CollectContainerEvents)
+
+	if cm.config.FinalizeTime > 0 && cm.config.FinalizeTime > cm.config.Interval {
+		// Add a timer for finalizing the application profile
+		startContainerTimer(id, cm.config.FinalizeTime, cm.FinalizeApplicationProfile)
+	}
 }
 
 func (cm *CollectorManager) ContainerStopped(id *ContainerId) {
@@ -352,6 +359,19 @@ func (cm *CollectorManager) CollectContainerEvents(id *ContainerId) {
 	}
 }
 
+func (cm *CollectorManager) FinalizeApplicationProfile(id *ContainerId) {
+	// Check if container is still running (is it in the map?)
+	if _, ok := cm.containers[*id]; ok {
+		// Patch the application profile to make it immutable with the final annotation
+		appProfileName := fmt.Sprintf("pod-%s", id.PodName)
+		_, err := cm.dynamicClient.Resource(AppProfileGvr).Namespace(id.Namespace).Patch(context.Background(),
+			appProfileName, apitypes.MergePatchType, []byte("{\"metadata\":{\"annotations\":{\"kapprofiler.kubescape.com/final\":\"true\"}}}"), v1.PatchOptions{})
+		if err != nil {
+			log.Printf("error patching application profile: %s\n", err)
+		}
+	}
+}
+
 func (cm *CollectorManager) doesApplicationProfileExists(namespace string, podName string, checkFinal bool, checkOwner bool) (bool, error) {
 	workloadKind := "Pod"
 	workloadName := podName

diff --git a/pkg/controller/controller.go b/pkg/controller/controller.go
@@ -108,7 +108,8 @@ func (c *Controller) handleApplicationProfile(obj interface{}) {
 						APIVersion: collector.ApplicationProfileApiVersion,
 					},
 					ObjectMeta: metav1.ObjectMeta{
-						Name: profileName,
+						Name:        profileName,
+						Annotations: applicationProfile.GetAnnotations(),
 					},
 					Spec: collector.ApplicationProfileSpec{
 						Containers: applicationProfile.Spec.Containers,
@@ -130,6 +131,7 @@ func (c *Controller) handleApplicationProfile(obj interface{}) {
 				}
 
 				deploymentApplicationProfile := &collector.ApplicationProfile{}
+				deploymentApplicationProfile.Annotations = applicationProfile.GetAnnotations()
 				deploymentApplicationProfile.Spec.Containers = applicationProfile.Spec.Containers
 				deploymentApplicationProfileRaw, _ := json.Marshal(deploymentApplicationProfile)
 				_, err = c.dynamicClient.Resource(collector.AppProfileGvr).Namespace(replicaSet.Namespace).Patch(context.TODO(), profileName, apitypes.MergePatchType, deploymentApplicationProfileRaw, metav1.PatchOptions{})
@@ -276,7 +278,8 @@ func (c *Controller) handleApplicationProfile(obj interface{}) {
 				APIVersion: collector.ApplicationProfileApiVersion,
 			},
 			ObjectMeta: metav1.ObjectMeta{
-				Name: applicationProfileNameForController,
+				Name:        applicationProfileNameForController,
+				Annotations: applicationProfile.GetAnnotations(),
 			},
 			Spec: collector.ApplicationProfileSpec{
 				Containers: containers,
@@ -299,6 +302,7 @@ func (c *Controller) handleApplicationProfile(obj interface{}) {
 			return
 		}
 		controllerApplicationProfile := &collector.ApplicationProfile{}
+		controllerApplicationProfile.Annotations = applicationProfile.GetAnnotations()
 		controllerApplicationProfile.Spec.Containers = containers
 		controllerApplicationProfileRaw, _ := json.Marshal(controllerApplicationProfile)
 		_, err = c.dynamicClient.Resource(collector.AppProfileGvr).Namespace(pod.Namespace).Patch(context.TODO(), applicationProfileNameForController, apitypes.MergePatchType, controllerApplicationProfileRaw, metav1.PatchOptions{})