From 5b4eb7f75cf4957764bdb5dc4a016d1692fe68b5 Mon Sep 17 00:00:00 2001
From: Tamal Saha <tamal@appscode.com>
Date: Fri, 6 Dec 2024 20:11:46 -0800
Subject: [PATCH] Add CheckImageArchitectures

Signed-off-by: Tamal Saha <tamal@appscode.com>
---
 pkg/lib/tests.go | 69 ++++++++++++++++++++++++++++++++++++++++++++++++
 pkg/lib/trivy.go | 12 ++++-----
 2 files changed, 75 insertions(+), 6 deletions(-)

diff --git a/pkg/lib/tests.go b/pkg/lib/tests.go
index 2caaa53a..a4bea392 100644
--- a/pkg/lib/tests.go
+++ b/pkg/lib/tests.go
@@ -17,10 +17,12 @@ limitations under the License.
 package lib
 
 import (
+	"errors"
 	"fmt"
 	"os"
 	"strings"
 
+	v1 "github.com/google/go-containerregistry/pkg/v1"
 	"k8s.io/apimachinery/pkg/util/sets"
 	"sigs.k8s.io/yaml"
 )
@@ -51,6 +53,73 @@ func CheckImageExists(files []string) error {
 	return nil
 }
 
+var desiredArchs = sets.New("amd64", "arm64")
+
+func CheckImageArchitectures(files []string, archSkipList []string) error {
+	archSkipSet := sets.NewString(archSkipList...)
+
+	images, err := LoadImageList(files)
+	if err != nil {
+		return err
+	}
+
+	var missing []string
+	missingArchs := map[string][]string{}
+	for _, img := range images {
+		obj, found, err := ImageManifest(img)
+		if err != nil || !found {
+			missing = append(missing, img)
+			continue
+		}
+		switch mf := obj.(type) {
+		case *v1.IndexManifest:
+			var archs []string
+			for _, d := range mf.Manifests {
+				if d.Platform != nil && d.Platform.Architecture != "" {
+					archs = append(archs, d.Platform.Architecture)
+				}
+			}
+			if missing := sets.List(desiredArchs.Difference(sets.New[string](archs...))); len(missing) > 0 {
+				missingArchs[img] = missing
+			} else {
+				fmt.Println("✔ " + img)
+			}
+		case *v1.Manifest:
+			if mf.Config.MediaType != "application/vnd.cncf.helm.config.v1+json" {
+				missingArchs[img] = []string{"arm64"}
+			}
+		default:
+			missingArchs[img] = []string{"amd64", "arm64"}
+		}
+	}
+
+	var fail bool
+	if len(missing) > 0 {
+		fmt.Println("----------------------------------------")
+		fmt.Println("Missing Images:")
+		fmt.Println(strings.Join(missing, "\n"))
+		fail = true
+	}
+
+	if len(missingArchs) > 0 {
+		fmt.Println("----------------------------------------")
+		fmt.Println("Missing Architectures:")
+		for img, archs := range missingArchs {
+			if !archSkipSet.Has(img) {
+				fmt.Printf("X %s %v\n", img, archs)
+				fail = true
+			} else {
+				fmt.Printf("[skipped] %s %v\n", img, archs)
+			}
+		}
+	}
+
+	if fail {
+		return errors.New("missing images and/or architectures")
+	}
+	return nil
+}
+
 func LoadImageList(files []string) ([]string, error) {
 	result := sets.New[string]()
 	for _, filename := range files {
diff --git a/pkg/lib/trivy.go b/pkg/lib/trivy.go
index 1bda90f2..78b58255 100644
--- a/pkg/lib/trivy.go
+++ b/pkg/lib/trivy.go
@@ -17,6 +17,7 @@ limitations under the License.
 package lib
 
 import (
+	"encoding/json"
 	"errors"
 	"fmt"
 	"net/http"
@@ -27,7 +28,6 @@ import (
 	"github.com/google/go-containerregistry/pkg/v1/remote/transport"
 	shell "gomodules.xyz/go-sh"
 	"kubeops.dev/scanner/apis/trivy"
-	"sigs.k8s.io/yaml"
 )
 
 // trivy image ubuntu --security-checks vuln --format json --quiet
@@ -80,28 +80,28 @@ func ImageManifest(ref string) (any, bool, error) {
 	data, err := crane.Manifest(ref, crane.WithAuthFromKeychain(authn.DefaultKeychain))
 	if err != nil {
 		if ImageNotFound(err) {
-			return "", false, nil
+			return nil, false, nil
 		}
 		return nil, false, err
 	}
 	var obj map[string]any
-	if err = yaml.Unmarshal(data, &obj); err != nil {
+	if err = json.Unmarshal(data, &obj); err != nil {
 		return nil, false, err
 	}
 	if _, ok := obj["manifests"]; ok {
 		var mf v1.IndexManifest
-		if err := yaml.Unmarshal(data, &mf); err != nil {
+		if err := json.Unmarshal(data, &mf); err != nil {
 			return nil, false, err
 		}
 		return &mf, true, nil
 	} else if _, ok := obj["layers"]; ok {
 		var mf v1.Manifest
-		if err := yaml.Unmarshal(data, &mf); err != nil {
+		if err := json.Unmarshal(data, &mf); err != nil {
 			return nil, false, err
 		}
 		return &mf, true, nil
 	}
-	return nil, true, fmt.Errorf("unknown image manifest format")
+	return nil, false, fmt.Errorf("unknown image manifest format")
 }
 
 func ImageNotFound(err error) bool {