From a905204e1c04e2f83caf10fd32091c86dede8cc1 Mon Sep 17 00:00:00 2001
From: RachelElysia <71795832+RachelElysia@users.noreply.github.com>
Date: Tue, 9 Jul 2024 09:14:25 -0400
Subject: [PATCH] Fleetctl: Add YAML validation error for apply command
 (#20254)

---
 changes/20075-fleetctl-apply-validation |  1 +
 cmd/fleetctl/apply.go                   | 11 ++++
 cmd/fleetctl/apply_test.go              | 80 ++++++++++++++++++++++++-
 3 files changed, 91 insertions(+), 1 deletion(-)
 create mode 100644 changes/20075-fleetctl-apply-validation

diff --git a/changes/20075-fleetctl-apply-validation b/changes/20075-fleetctl-apply-validation
new file mode 100644
index 0000000000..02831f8b54
--- /dev/null
+++ b/changes/20075-fleetctl-apply-validation
@@ -0,0 +1 @@
+- Add .yml and .yaml file type validation and error message to fleetctl apply
diff --git a/cmd/fleetctl/apply.go b/cmd/fleetctl/apply.go
index 103f7515e1..751c74709f 100644
--- a/cmd/fleetctl/apply.go
+++ b/cmd/fleetctl/apply.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
+	"strings"
 
 	"github.com/fleetdm/fleet/v4/pkg/spec"
 	"github.com/fleetdm/fleet/v4/server/fleet"
@@ -61,6 +62,16 @@ func applyCommand() *cli.Command {
 			if err != nil {
 				return err
 			}
+
+			// Check if the file has a .yml or .yaml extension
+			ext := strings.ToLower(filepath.Ext(flFilename))
+			if ext == "" {
+				return errors.New("Missing file extension: only .yml or .yaml files can be applied")
+			}
+			if ext != ".yml" && ext != ".yaml" {
+				return fmt.Errorf("Invalid file extension %s: only .yml or .yaml files can be applied", ext)
+			}
+
 			specs, err := spec.GroupFromBytes(b)
 			if err != nil {
 				return err
diff --git a/cmd/fleetctl/apply_test.go b/cmd/fleetctl/apply_test.go
index b9a741742a..5798e26aa8 100644
--- a/cmd/fleetctl/apply_test.go
+++ b/cmd/fleetctl/apply_test.go
@@ -31,6 +31,7 @@ import (
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"github.com/urfave/cli/v2"
 )
 
 var userRoleSpecList = []*fleet.User{
@@ -1672,7 +1673,6 @@ func TestApplyLabels(t *testing.T) {
 	_, err = runAppNoChecks([]string{"apply", "-f", name})
 	require.Error(t, err)
 	require.ErrorContains(t, err, "cannot modify or add built-in label")
-
 }
 
 func TestApplyPacks(t *testing.T) {
@@ -3760,3 +3760,81 @@ spec:
 		})
 	}
 }
+
+func TestApplyFileExtensionValidation(t *testing.T) {
+	cases := []struct {
+		desc     string
+		filename string
+		wantErr  string
+	}{
+		{
+			desc:     "Valid .yml extension",
+			filename: "test_file.yml",
+			wantErr:  "",
+		},
+		{
+			desc:     "Valid .yaml extension",
+			filename: "test_file.yaml",
+			wantErr:  "",
+		},
+		{
+			desc:     "Invalid .txt extension",
+			filename: "test_file.txt",
+			wantErr:  "Invalid file extension .txt: only .yml or .yaml files can be applied",
+		},
+		{
+			desc:     "Invalid .json extension",
+			filename: "test_file.json",
+			wantErr:  "Invalid file extension .json: only .yml or .yaml files can be applied",
+		},
+		{
+			desc:     "No extension",
+			filename: "test_file",
+			wantErr:  "Missing file extension: only .yml or .yaml files can be applied",
+		},
+	}
+
+	for _, c := range cases {
+		t.Run(c.desc, func(t *testing.T) {
+			// Create a temporary directory
+			tmpDir, err := os.MkdirTemp("", "test")
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer os.RemoveAll(tmpDir) // clean up
+
+			// Create the file with the exact name in the temporary directory
+			tmpFilePath := filepath.Join(tmpDir, c.filename)
+			tmpFile, err := os.Create(tmpFilePath)
+			if err != nil {
+				t.Fatal(err)
+			}
+			tmpFile.Close()
+
+			// Create a new cli.App for each test
+			app := &cli.App{
+				Commands: []*cli.Command{
+					applyCommand(),
+				},
+			}
+
+			// Set up arguments
+			args := []string{"fleetctl", "apply", "-f", tmpFilePath}
+
+			// Run the command
+			err = app.Run(args)
+
+			if c.wantErr == "" {
+				if err != nil {
+					t.Errorf("Expected no error, but got: %v", err)
+				}
+			} else {
+				if err == nil {
+					t.Errorf("Expected error, but got none")
+				} else if err.Error() != c.wantErr {
+					t.Errorf("Expected error message '%s', but got '%s'", c.wantErr, err.Error())
+				}
+			}
+		})
+	}
+}