From 0962ace1d1d87390ce80511b3f2c710a21429eda Mon Sep 17 00:00:00 2001
From: mandochen <mandochen@tencent.com>
Date: Mon, 5 Jun 2023 20:40:41 +0800
Subject: [PATCH] bugfix: handle default images register

---
 cmd/main.go       |  2 +-
 kubectl_images.go | 24 +++++++++++++++++++++---
 2 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/cmd/main.go b/cmd/main.go
index 6f6a525..0c84b8d 100644
--- a/cmd/main.go
+++ b/cmd/main.go
@@ -9,7 +9,7 @@ import (
 	"github.com/spf13/cobra"
 )
 
-const version = "0.6.1"
+const version = "0.6.2"
 
 var rootCmd *cobra.Command
 
diff --git a/kubectl_images.go b/kubectl_images.go
index e1926f4..5e5227d 100644
--- a/kubectl_images.go
+++ b/kubectl_images.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"os"
 	"os/exec"
+	"path"
 	"regexp"
 	"strconv"
 	"strings"
@@ -191,6 +192,11 @@ func (ki *KubeImage) nodeCommands() []string {
 	return append([]string{"get", "nodes", "-o", nodeTemplate}, kubecfg...)
 }
 
+func (ki *KubeImage) recordImageSize(image string, size int) {
+	ki.imageSize[image] = size
+	ki.imageSize[path.Base(image)] = size
+}
+
 func (ki *KubeImage) execNodeCommand() {
 	process := exec.Command("kubectl", ki.nodeCommands()...)
 	bs, err := process.CombinedOutput()
@@ -202,17 +208,29 @@ func (ki *KubeImage) execNodeCommand() {
 	for _, line := range stringSplit(string(bs), "\n") {
 		items := stringSplit(line, ",")
 		switch len(items) {
+		case 2:
+			size, err := strconv.Atoi(items[1])
+			if err != nil {
+				continue
+			}
+
+			ki.recordImageSize(items[0], size)
+			parts := strings.Split(items[0], ":")
+			if len(parts) == 2 && parts[1] == "latest" {
+				ki.recordImageSize(parts[0], size)
+			}
+
 		case 3:
 			size, err := strconv.Atoi(items[2])
 			if err != nil {
 				continue
 			}
 
-			ki.imageSize[items[0]] = size
-			ki.imageSize[items[1]] = size
+			ki.recordImageSize(items[0], size)
+			ki.recordImageSize(items[1], size)
 			parts := strings.Split(items[1], ":")
 			if len(parts) == 2 && parts[1] == "latest" {
-				ki.imageSize[parts[0]] = size
+				ki.recordImageSize(parts[0], size)
 			}
 		}
 	}