From 6528f3ed648f05903b11a061741748a41e0f4c7e Mon Sep 17 00:00:00 2001
From: phm07 <22707808+phm07@users.noreply.github.com>
Date: Wed, 14 Aug 2024 13:37:35 +0200
Subject: [PATCH] fix(image): only accept numerical ids for update, delete,
 label

---
 internal/cmd/image/delete.go      | 11 +++++++++--
 internal/cmd/image/delete_test.go | 29 +++++++++++++----------------
 internal/cmd/image/describe.go    |  9 +++++++--
 internal/cmd/image/labels.go      |  7 ++++++-
 internal/cmd/image/labels_test.go |  4 ++--
 internal/cmd/image/update.go      | 11 +++++++++--
 internal/cmd/image/update_test.go |  4 ++--
 7 files changed, 48 insertions(+), 27 deletions(-)

diff --git a/internal/cmd/image/delete.go b/internal/cmd/image/delete.go
index 77066e76..aabe5649 100644
--- a/internal/cmd/image/delete.go
+++ b/internal/cmd/image/delete.go
@@ -1,6 +1,9 @@
 package image
 
 import (
+	"fmt"
+	"strconv"
+
 	"github.com/spf13/cobra"
 
 	"github.com/hetznercloud/cli/internal/cmd/base"
@@ -14,8 +17,12 @@ var DeleteCmd = base.DeleteCmd{
 	ResourceNamePlural:   "images",
 	ShortDescription:     "Delete an image",
 	NameSuggestions:      func(c hcapi2.Client) func() []string { return c.Image().Names },
-	Fetch: func(s state.State, cmd *cobra.Command, idOrName string) (interface{}, *hcloud.Response, error) {
-		return s.Client().Image().Get(s, idOrName)
+	Fetch: func(s state.State, _ *cobra.Command, idOrName string) (interface{}, *hcloud.Response, error) {
+		id, err := strconv.ParseInt(idOrName, 10, 64)
+		if err != nil {
+			return nil, nil, fmt.Errorf("invalid snapshot or backup ID %q", idOrName)
+		}
+		return s.Client().Image().GetByID(s, id)
 	},
 	Delete: func(s state.State, cmd *cobra.Command, resource interface{}) (*hcloud.Action, error) {
 		image := resource.(*hcloud.Image)
diff --git a/internal/cmd/image/delete_test.go b/internal/cmd/image/delete_test.go
index 59d0a1e7..ef40f5d4 100644
--- a/internal/cmd/image/delete_test.go
+++ b/internal/cmd/image/delete_test.go
@@ -1,6 +1,7 @@
 package image_test
 
 import (
+	"strconv"
 	"testing"
 
 	"github.com/golang/mock/gomock"
@@ -19,20 +20,19 @@ func TestDelete(t *testing.T) {
 	fx.ExpectEnsureToken()
 
 	img := &hcloud.Image{
-		ID:   123,
-		Name: "test",
+		ID: 123,
 	}
 
 	fx.Client.ImageClient.EXPECT().
-		Get(gomock.Any(), "test").
+		GetByID(gomock.Any(), img.ID).
 		Return(img, nil, nil)
 	fx.Client.ImageClient.EXPECT().
 		Delete(gomock.Any(), img).
 		Return(nil, nil)
 
-	out, errOut, err := fx.Run(cmd, []string{"test"})
+	out, errOut, err := fx.Run(cmd, []string{"123"})
 
-	expOut := "image test deleted\n"
+	expOut := "image 123 deleted\n"
 
 	assert.NoError(t, err)
 	assert.Empty(t, errOut)
@@ -48,33 +48,30 @@ func TestDeleteMultiple(t *testing.T) {
 
 	images := []*hcloud.Image{
 		{
-			ID:   123,
-			Name: "test1",
+			ID: 123,
 		},
 		{
-			ID:   456,
-			Name: "test2",
+			ID: 456,
 		},
 		{
-			ID:   789,
-			Name: "test3",
+			ID: 789,
 		},
 	}
 
-	var names []string
+	var ids []string
 	for _, img := range images {
-		names = append(names, img.Name)
+		ids = append(ids, strconv.FormatInt(img.ID, 10))
 		fx.Client.ImageClient.EXPECT().
-			Get(gomock.Any(), img.Name).
+			GetByID(gomock.Any(), img.ID).
 			Return(img, nil, nil)
 		fx.Client.ImageClient.EXPECT().
 			Delete(gomock.Any(), img).
 			Return(nil, nil)
 	}
 
-	out, errOut, err := fx.Run(cmd, names)
+	out, errOut, err := fx.Run(cmd, ids)
 
 	assert.NoError(t, err)
 	assert.Empty(t, errOut)
-	assert.Equal(t, "images test1, test2, test3 deleted\n", out)
+	assert.Equal(t, "images 123, 456, 789 deleted\n", out)
 }
diff --git a/internal/cmd/image/describe.go b/internal/cmd/image/describe.go
index 327e22fa..22d8138b 100644
--- a/internal/cmd/image/describe.go
+++ b/internal/cmd/image/describe.go
@@ -3,6 +3,7 @@ package image
 import (
 	"fmt"
 	"os"
+	"strconv"
 
 	humanize "github.com/dustin/go-humanize"
 	"github.com/spf13/cobra"
@@ -23,17 +24,21 @@ var DescribeCmd = base.DescribeCmd{
 	AdditionalFlags: func(cmd *cobra.Command) {
 		cmd.Flags().StringP("architecture", "a", string(hcloud.ArchitectureX86), "architecture of the image, default is x86")
 		cmd.RegisterFlagCompletionFunc("architecture", cmpl.SuggestCandidates(string(hcloud.ArchitectureX86), string(hcloud.ArchitectureARM)))
-
 	},
 	NameSuggestions: func(c hcapi2.Client) func() []string { return c.Image().Names },
 	Fetch: func(s state.State, cmd *cobra.Command, idOrName string) (interface{}, interface{}, error) {
+		_, err := strconv.ParseInt(idOrName, 10, 64)
+		isId := err == nil
+
 		arch, err := cmd.Flags().GetString("architecture")
 		if err != nil {
 			return nil, nil, err
 		}
-		if !cmd.Flags().Changed("architecture") {
+
+		if !isId && !cmd.Flags().Changed("architecture") {
 			_, _ = fmt.Fprintln(os.Stderr, "INFO: This command only returns x86 images by default. Explicitly set the --architecture=x86|arm flag to hide this message.")
 		}
+
 		img, _, err := s.Client().Image().GetForArchitecture(s, idOrName, hcloud.Architecture(arch))
 		if err != nil {
 			return nil, nil, err
diff --git a/internal/cmd/image/labels.go b/internal/cmd/image/labels.go
index e1ae3849..1c4f4640 100644
--- a/internal/cmd/image/labels.go
+++ b/internal/cmd/image/labels.go
@@ -2,6 +2,7 @@ package image
 
 import (
 	"fmt"
+	"strconv"
 
 	"github.com/hetznercloud/cli/internal/cmd/base"
 	"github.com/hetznercloud/cli/internal/hcapi2"
@@ -16,7 +17,11 @@ var LabelCmds = base.LabelCmds{
 	NameSuggestions:        func(c hcapi2.Client) func() []string { return c.Image().Names },
 	LabelKeySuggestions:    func(c hcapi2.Client) func(idOrName string) []string { return c.Image().LabelKeys },
 	FetchLabels: func(s state.State, idOrName string) (map[string]string, int64, error) {
-		image, _, err := s.Client().Image().Get(s, idOrName)
+		id, err := strconv.ParseInt(idOrName, 10, 64)
+		if err != nil {
+			return nil, 0, fmt.Errorf("invalid snapshot or backup ID %q", idOrName)
+		}
+		image, _, err := s.Client().Image().GetByID(s, id)
 		if err != nil {
 			return nil, 0, err
 		}
diff --git a/internal/cmd/image/labels_test.go b/internal/cmd/image/labels_test.go
index 8a10839b..47c5d0e7 100644
--- a/internal/cmd/image/labels_test.go
+++ b/internal/cmd/image/labels_test.go
@@ -19,7 +19,7 @@ func TestLabelAdd(t *testing.T) {
 	fx.ExpectEnsureToken()
 
 	fx.Client.ImageClient.EXPECT().
-		Get(gomock.Any(), "123").
+		GetByID(gomock.Any(), int64(123)).
 		Return(&hcloud.Image{ID: 123}, nil, nil)
 	fx.Client.ImageClient.EXPECT().
 		Update(gomock.Any(), &hcloud.Image{ID: 123}, hcloud.ImageUpdateOpts{
@@ -45,7 +45,7 @@ func TestLabelRemove(t *testing.T) {
 	fx.ExpectEnsureToken()
 
 	fx.Client.ImageClient.EXPECT().
-		Get(gomock.Any(), "123").
+		GetByID(gomock.Any(), int64(123)).
 		Return(&hcloud.Image{
 			ID: 123,
 			Labels: map[string]string{
diff --git a/internal/cmd/image/update.go b/internal/cmd/image/update.go
index 75e2562f..5aaa5e5b 100644
--- a/internal/cmd/image/update.go
+++ b/internal/cmd/image/update.go
@@ -1,6 +1,9 @@
 package image
 
 import (
+	"fmt"
+	"strconv"
+
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
 
@@ -15,8 +18,12 @@ var UpdateCmd = base.UpdateCmd{
 	ResourceNameSingular: "Image",
 	ShortDescription:     "Update an image",
 	NameSuggestions:      func(c hcapi2.Client) func() []string { return c.Image().Names },
-	Fetch: func(s state.State, cmd *cobra.Command, idOrName string) (interface{}, *hcloud.Response, error) {
-		return s.Client().Image().Get(s, idOrName)
+	Fetch: func(s state.State, _ *cobra.Command, idOrName string) (interface{}, *hcloud.Response, error) {
+		id, err := strconv.ParseInt(idOrName, 10, 64)
+		if err != nil {
+			return nil, nil, fmt.Errorf("invalid snapshot or backup ID %q", idOrName)
+		}
+		return s.Client().Image().GetByID(s, id)
 	},
 	DefineFlags: func(cmd *cobra.Command) {
 		cmd.Flags().String("description", "", "Image description")
diff --git a/internal/cmd/image/update_test.go b/internal/cmd/image/update_test.go
index 44a5821a..978a212e 100644
--- a/internal/cmd/image/update_test.go
+++ b/internal/cmd/image/update_test.go
@@ -19,7 +19,7 @@ func TestUpdateDescription(t *testing.T) {
 	fx.ExpectEnsureToken()
 
 	fx.Client.ImageClient.EXPECT().
-		Get(gomock.Any(), "123").
+		GetByID(gomock.Any(), int64(123)).
 		Return(&hcloud.Image{ID: 123}, nil, nil)
 	fx.Client.ImageClient.EXPECT().
 		Update(gomock.Any(), &hcloud.Image{ID: 123}, hcloud.ImageUpdateOpts{
@@ -43,7 +43,7 @@ func TestUpdateType(t *testing.T) {
 	fx.ExpectEnsureToken()
 
 	fx.Client.ImageClient.EXPECT().
-		Get(gomock.Any(), "123").
+		GetByID(gomock.Any(), int64(123)).
 		Return(&hcloud.Image{ID: 123}, nil, nil)
 	fx.Client.ImageClient.EXPECT().
 		Update(gomock.Any(), &hcloud.Image{ID: 123}, hcloud.ImageUpdateOpts{