From 594c51ced32b7df2cfacdb37a781901a876bcbaa Mon Sep 17 00:00:00 2001
From: Bishwa Thapa <thapabishwakranti+dev@gmail.com>
Date: Tue, 29 Aug 2023 23:55:16 +0545
Subject: [PATCH] feat: enhancement to list workload if owner is a custom
 resource

---
 examples/trivy.go        | 11 +++++++++--
 pkg/trivyk8s/trivyk8s.go | 18 +++++++++++++++---
 2 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/examples/trivy.go b/examples/trivy.go
index 80330b3..52b1a2a 100644
--- a/examples/trivy.go
+++ b/examples/trivy.go
@@ -30,12 +30,19 @@ func main() {
 
 	fmt.Println("Current namespace:", cluster.GetCurrentNamespace())
 
-	trivyk8s := trivyk8s.New(cluster, logger.Sugar())
+	trivyk8s := trivyk8s.New(cluster, logger.Sugar(), trivyk8s.WithExcludeOwned(true))
+
+	fmt.Println("Scanning kind 'pods' with exclude-owned=true")
+	artifacts, err := trivyk8s.Resources("pod").AllNamespaces().ListArtifacts(ctx)
+	if err != nil {
+		log.Fatal(err)
+	}
+	printArtifacts(artifacts)
 
 	fmt.Println("Scanning cluster")
 
 	//trivy k8s #cluster
-	artifacts, err := trivyk8s.ListArtifacts(ctx)
+	artifacts, err = trivyk8s.ListArtifacts(ctx)
 	if err != nil {
 		log.Fatal(err)
 	}
diff --git a/pkg/trivyk8s/trivyk8s.go b/pkg/trivyk8s/trivyk8s.go
index bd3ddd9..aa98875 100644
--- a/pkg/trivyk8s/trivyk8s.go
+++ b/pkg/trivyk8s/trivyk8s.go
@@ -139,9 +139,21 @@ func (c *client) ListArtifacts(ctx context.Context) ([]*artifacts.Artifact, erro
 				continue
 			}
 
-			// if excluding owned resources is enabled, we check if the resource has an owner
-			// if it does, then skip it
-			if c.excludeOwned && len(resource.GetOwnerReferences()) > 0 {
+			// assume that the owner is a built-in workload by default
+			ownerIsBuiltIn := true
+			if len(resource.GetOwnerReferences()) > 0 {
+				// if the resource has an owner, we check if it is a built-in workload
+				// this ensures that we don't skip resources that are owned by custom resources
+				for _, owner := range resource.GetOwnerReferences() {
+					if !k8s.IsBuiltInWorkload(&owner) {
+						ownerIsBuiltIn = false
+						break
+					}
+				}
+			}
+
+			// if excludeOwned is enabled and workload is a built-in workload, we skip it
+			if c.excludeOwned && ownerIsBuiltIn {
 				continue
 			}