CyberArk: Get --input-file working with machinehub mode

api/datareading.go

@@ -1,9 +1,12 @@
 package api
 
 import (
+	"bytes"
 	"encoding/json"
+	"fmt"
 	"time"
 
+	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	"k8s.io/apimachinery/pkg/version"
 )
 
@@ -26,12 +29,64 @@ type DataReading struct {
 	SchemaVersion string      `json:"schema_version"`
 }
 
+// UnmarshalJSON implements the json.Unmarshaler interface for DataReading.
+// It handles the dynamic parsing of the Data field based on the DataGatherer.
+func (o *DataReading) UnmarshalJSON(data []byte) error {
+	var tmp struct {
+		ClusterID     string          `json:"cluster_id,omitempty"`
+		DataGatherer  string          `json:"data-gatherer"`
+		Timestamp     Time            `json:"timestamp"`
+		Data          json.RawMessage `json:"data"`
+		SchemaVersion string          `json:"schema_version"`
+	}
+
+	d := json.NewDecoder(bytes.NewReader(data))
+	d.DisallowUnknownFields()
+
+	if err := d.Decode(&tmp); err != nil {
+		return err
+	}
+	o.ClusterID = tmp.ClusterID
+	o.DataGatherer = tmp.DataGatherer
+	o.Timestamp = tmp.Timestamp
+	o.SchemaVersion = tmp.SchemaVersion
+
+	{
+		var discoveryData DiscoveryData
+		d := json.NewDecoder(bytes.NewReader(tmp.Data))
+		d.DisallowUnknownFields()
+		if err := d.Decode(&discoveryData); err == nil {
+			o.Data = &discoveryData
+			return nil
+		}
+	}
+	{
+		var dynamicData DynamicData
+		d := json.NewDecoder(bytes.NewReader(tmp.Data))
+		d.DisallowUnknownFields()
+		if err := d.Decode(&dynamicData); err == nil {
+			o.Data = &dynamicData
+			return nil
+		}
+	}
+	{
+		var genericData map[string]interface{}
+		d := json.NewDecoder(bytes.NewReader(tmp.Data))
+		d.DisallowUnknownFields()
+		if err := d.Decode(&genericData); err == nil {
+			o.Data = genericData
+			return nil
+		}
+	}
+	return fmt.Errorf("failed to parse DataReading.Data for gatherer %s", o.DataGatherer)
+}
+
 // GatheredResource wraps the raw k8s resource that is sent to the jetstack secure backend
 type GatheredResource struct {
 	// Resource is a reference to a k8s object that was found by the informer
 	// should be of type unstructured.Unstructured, raw Object
-	Resource  interface{}
-	DeletedAt Time
+	Resource  interface{} `json:"resource"`
+	DeletedAt Time        `json:"deleted_at,omitempty"`
 }
 
 func (v GatheredResource) MarshalJSON() ([]byte, error) {
@@ -51,6 +106,23 @@ func (v GatheredResource) MarshalJSON() ([]byte, error) {
 	return json.Marshal(data)
 }
 
+func (v *GatheredResource) UnmarshalJSON(data []byte) error {
+	var tmpResource struct {
+		Resource  *unstructured.Unstructured `json:"resource"`
+		DeletedAt Time                       `json:"deleted_at,omitempty"`
+	}
+
+	d := json.NewDecoder(bytes.NewReader(data))
+	d.DisallowUnknownFields()
+
+	if err := d.Decode(&tmpResource); err != nil {
+		return err
+	}
+	v.Resource = tmpResource.Resource
+	v.DeletedAt = tmpResource.DeletedAt
+	return nil
+}
+
 // DynamicData is the DataReading.Data returned by the k8s.DataGathererDynamic
 // gatherer
 type DynamicData struct {

cmd/agent_test.go

@@ -3,49 +3,81 @@ package cmd
 import (
 	"bytes"
 	"context"
+	"fmt"
 	"os"
 	"os/exec"
+	"path/filepath"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/require"
+
+	arktesting "github.com/jetstack/preflight/internal/cyberark/testing"
 )
 
-// TestAgentRunOneShot runs the agent in `--one-shot` mode and verifies that it exits
-// after the first data gathering iteration.
-func TestAgentRunOneShot(t *testing.T) {
+// TestOutputModes tests the different output modes of the agent command.
+// It does this by running the agent command in a subprocess with the
+// appropriate flags and configuration files.
+// It assumes that the test is being run from the "cmd" directory and that
+// the repository root is the parent directory of the current working directory.
+func TestOutputModes(t *testing.T) {
+	repoRoot := findRepoRoot(t)
+
+	t.Run("localfile", func(t *testing.T) {
+		runSubprocess(t, repoRoot, []string{
+			"--agent-config-file", filepath.Join(repoRoot, "examples/localfile/config.yaml"),
+			"--input-path", filepath.Join(repoRoot, "examples/localfile/input.json"),
+			"--output-path", "/dev/null",
+		})
+	})
+
+	t.Run("machinehub", func(t *testing.T) {
+		arktesting.SkipIfNoEnv(t)
+		runSubprocess(t, repoRoot, []string{
+			"--agent-config-file", filepath.Join(repoRoot, "examples/machinehub/config.yaml"),
+			"--input-path", filepath.Join(repoRoot, "examples/machinehub/input.json"),
+			"--machine-hub",
+		})
+	})
+}
+
+// findRepoRoot returns the absolute path to the repository root.
+// It assumes that the test is being run from the "cmd" directory.
+func findRepoRoot(t *testing.T) string {
+	cwd, err := os.Getwd()
+	require.NoError(t, err)
+	repoRoot, err := filepath.Abs(filepath.Join(cwd, ".."))
+	require.NoError(t, err)
+	return repoRoot
+}
+
+// runSubprocess runs the current test in a subprocess with the given args.
+// It sets the GO_CHILD environment variable to indicate to the subprocess
+// that it should run the main function instead of the test function.
+// It captures and logs the stdout and stderr of the subprocess.
+// It fails the test if the subprocess exits with a non-zero status.
+// It uses a timeout to avoid hanging indefinitely.
+func runSubprocess(t *testing.T, repoRoot string, args []string) {
 	if _, found := os.LookupEnv("GO_CHILD"); found {
-		os.Args = []string{
+		os.Args = append([]string{
 			"preflight",
 			"agent",
+			"--log-level", "6",
 			"--one-shot",
-			"--agent-config-file=testdata/agent/one-shot/success/config.yaml",
-			"--input-path=testdata/agent/one-shot/success/input.json",
-			"--output-path=/dev/null",
-			"-v=9",
-		}
+		}, args...)
 		Execute()
 		return
 	}
-	t.Log("Running child process")
-	ctx, cancel := context.WithTimeout(t.Context(), time.Second*3)
+	t.Log("Running child process", os.Args[0], "-test.run=^"+t.Name()+"$")
+	ctx, cancel := context.WithTimeout(t.Context(), time.Second*10)
 	defer cancel()
-	cmd := exec.CommandContext(ctx, os.Args[0], "-test.run=^TestAgentRunOneShot$")
-	var (
-		stdout bytes.Buffer
-		stderr bytes.Buffer
-	)
+	cmd := exec.CommandContext(ctx, os.Args[0], "-test.run=^"+t.Name()+"$")
+	var stdout, stderr bytes.Buffer
 	cmd.Stdout = &stdout
 	cmd.Stderr = &stderr
-	cmd.Env = append(
-		os.Environ(),
-		"GO_CHILD=true",
-	)
+	cmd.Env = append(os.Environ(), "GO_CHILD=true")
 	err := cmd.Run()
-
-	stdoutStr := stdout.String()
-	stderrStr := stderr.String()
-	t.Logf("STDOUT\n%s\n", stdoutStr)
-	t.Logf("STDERR\n%s\n", stderrStr)
-	require.NoError(t, err, context.Cause(ctx))
+	t.Logf("STDOUT\n%s\n", stdout.String())
+	t.Logf("STDERR\n%s\n", stderr.String())
+	require.NoError(t, err, fmt.Sprintf("Error: %v\nSTDERR: %s", err, stderr.String()))
 }

examples/localfile/config.yaml

@@ -0,0 +1 @@
+# No config is required to run the agent with an input file and an output file.

examples/machinehub/config.yaml

@@ -0,0 +1 @@
+# Not used

examples/machinehub/input.json

@@ -0,0 +1,105 @@
+[
+    {
+        "data-gatherer": "ark/discovery",
+        "data": {
+            "cluster_id": "0e069229-d83b-4075-a4c8-95838ff5c437",
+            "server_version": {
+                "gitVersion": "v1.27.6"
+            }
+        }
+    },
+    {
+        "data-gatherer": "ark/secrets",
+        "data": {
+            "items": [
+                {
+                    "resource": {
+                        "kind": "Secret",
+                        "apiVersion": "v1",
+                        "metadata": {
+                            "name": "app-1-secret-1",
+                            "namespace": "team-1"
+                        }
+                    }
+                }
+            ]
+        }
+    },
+    {
+        "data-gatherer": "ark/pods",
+        "data": {
+            "items": [
+                {
+                    "resource": {
+                        "kind": "Pod",
+                        "apiVersion": "v1",
+                        "metadata": {
+                            "name": "app-1-pod-1",
+                            "namespace": "team-1"
+                        }
+                    }
+                }
+            ]
+        }
+    },
+    {
+        "data-gatherer": "ark/statefulsets",
+        "data": {
+            "items": []
+        }
+    },
+    {
+        "data-gatherer": "ark/deployments",
+        "data": {
+            "items": []
+        }
+    },
+    {
+        "data-gatherer": "ark/clusterroles",
+        "data": {
+            "items": []
+        }
+    },
+    {
+        "data-gatherer": "ark/roles",
+        "data": {
+            "items": []
+        }
+    },
+    {
+        "data-gatherer": "ark/clusterrolebindings",
+        "data": {
+            "items": []
+        }
+    },
+    {
+        "data-gatherer": "ark/rolebindings",
+        "data": {
+            "items": []
+        }
+    },
+    {
+        "data-gatherer": "ark/cronjobs",
+        "data": {
+            "items": []
+        }
+    },
+    {
+        "data-gatherer": "ark/jobs",
+        "data": {
+            "items": []
+        }
+    },
+    {
+        "data-gatherer": "ark/daemonsets",
+        "data": {
+            "items": []
+        }
+    },
+    {
+        "data-gatherer": "ark/serviceaccounts",
+        "data": {
+            "items": []
+        }
+    }
+]