Simplify Kubernetes Injection

cmd/internal/cmderr/checks.go

@@ -0,0 +1,26 @@
+package cmderr
+
+import (
+	"errors"
+	"github.com/akitasoftware/akita-cli/cfg"
+	"github.com/akitasoftware/akita-cli/env"
+	"github.com/akitasoftware/akita-cli/printer"
+)
+
+// Checks that a user has configured their API key and secret and returned them.
+// If the user has not configured their API key, a user-friendly error message is printed and an error is returned.
+func RequireAPICredentials(explanation string) (string, string, error) {
+	key, secret := cfg.GetAPIKeyAndSecret()
+	if key == "" || secret == "" {
+		printer.Errorf("No Akita API key configured. %s\n", explanation)
+		if env.InDocker() {
+			printer.Infof("Please set the AKITA_API_KEY_ID and AKITA_API_KEY_SECRET environment variables on the Docker command line.\n")
+		} else {
+			printer.Infof("Use the AKITA_API_KEY_ID and AKITA_API_KEY_SECRET environment variables, or run 'akita login'.\n")
+		}
+
+		return "", "", AkitaErr{Err: errors.New("could not find an Akita API key to use")}
+	}
+
+	return key, secret, nil
+}

cmd/internal/ecs/ecs.go

@@ -4,10 +4,7 @@ import (
 	"fmt"
 	"strings"
 
-	"github.com/akitasoftware/akita-cli/cfg"
 	"github.com/akitasoftware/akita-cli/cmd/internal/cmderr"
-	"github.com/akitasoftware/akita-cli/env"
-	"github.com/akitasoftware/akita-cli/printer"
 	"github.com/akitasoftware/akita-cli/rest"
 	"github.com/akitasoftware/akita-cli/telemetry"
 	"github.com/akitasoftware/akita-cli/util"
@@ -71,8 +68,18 @@ func init() {
 	Cmd.PersistentFlags().StringVar(&awsRegionFlag, "region", "", "The AWS region in which your ECS cluster resides.")
 	Cmd.PersistentFlags().StringVar(&ecsClusterFlag, "cluster", "", "The name or ARN of your ECS cluster.")
 	Cmd.PersistentFlags().StringVar(&ecsServiceFlag, "service", "", "The name or ARN of your ECS service.")
-	Cmd.PersistentFlags().StringVar(&ecsTaskDefinitionFlag, "task", "", "The name of your ECS task definition to modify.")
-	Cmd.PersistentFlags().BoolVar(&dryRunFlag, "dry-run", false, "Perform a dry run: show what will be done, but do not modify ECS.")
+	Cmd.PersistentFlags().StringVar(
+		&ecsTaskDefinitionFlag,
+		"task",
+		"",
+		"The name of your ECS task definition to modify.",
+	)
+	Cmd.PersistentFlags().BoolVar(
+		&dryRunFlag,
+		"dry-run",
+		false,
+		"Perform a dry run: show what will be done, but do not modify ECS.",
+	)
 
 	// Support for credentials in a nonstandard location
 	Cmd.PersistentFlags().StringVar(&awsCredentialsFlag, "aws-credentials", "", "Location of AWS credentials file.")
@@ -84,29 +91,34 @@ func init() {
 
 func addAgentToECS(cmd *cobra.Command, args []string) error {
 	// Check for API key
-	key, secret := cfg.GetAPIKeyAndSecret()
-	if key == "" || secret == "" {
-		printer.Errorf("No Akita API key configured. The Akita agent must have an API key in order to capture traces.\n")
-		if env.InDocker() {
-			printer.Infof("Please set the AKITA_API_KEY_ID and AKITA_API_KEY_SECRET environment variables on the Docker command line.\n")
-		} else {
-			printer.Infof("Use the AKITA_API_KEY_ID and AKITA_API_KEY_SECRET environment variables, or run 'akita login'.\n")
-		}
-		return cmderr.AkitaErr{Err: errors.New("Could not find an Akita API key to use.")}
+	_, _, err := cmderr.RequireAPICredentials("The Akita agent must have an API key in order to capture traces.")
+	if err != nil {
+		return err
 	}
 
 	// Check project's existence
 	if projectFlag == "" {
 		return errors.New("Must specify the name of your Akita project with the --project flag.")
 	}
 	frontClient := rest.NewFrontClient(rest.Domain, telemetry.GetClientID())
-	_, err := util.GetServiceIDByName(frontClient, projectFlag)
+	_, err = util.GetServiceIDByName(frontClient, projectFlag)
 	if err != nil {
 		// TODO: we _could_ offer to create it, instead.
 		if strings.Contains(err.Error(), "cannot determine project ID") {
-			return cmderr.AkitaErr{Err: fmt.Errorf("Could not find the project %q in the Akita cloud. Please create it from the Akita web console before proceeding.", projectFlag)}
+			return cmderr.AkitaErr{
+				Err: fmt.Errorf(
+					"Could not find the project %q in the Akita cloud. Please create it from the Akita web console before proceeding.",
+					projectFlag,
+				),
+			}
 		} else {
-			return cmderr.AkitaErr{Err: errors.Wrapf(err, "Could not look up the project %q in the Akita cloud", projectFlag)}
+			return cmderr.AkitaErr{
+				Err: errors.Wrapf(
+					err,
+					"Could not look up the project %q in the Akita cloud",
+					projectFlag,
+				),
+			}
 		}
 	}
 

cmd/internal/kube/inject.go

@@ -0,0 +1,247 @@
+package kube
+
+import (
+	"bytes"
+
+	"github.com/akitasoftware/akita-cli/cmd/internal/cmderr"
+	"github.com/akitasoftware/akita-cli/cmd/internal/kube/injector"
+	"github.com/akitasoftware/akita-cli/printer"
+	"github.com/akitasoftware/akita-cli/telemetry"
+	"github.com/akitasoftware/go-utils/optionals"
+	"github.com/pkg/errors"
+	"github.com/spf13/cobra"
+	v1 "k8s.io/api/core/v1"
+)
+
+var (
+	// The target Yaml faile to be injected
+	// This is required for execution of injectCmd
+	injectFileNameFlag string
+	// The output file to write the injected Yaml to
+	// If not set, injectCmd will default to printing the output to stdout
+	injectOutputFlag string
+	// The name of the project that the injected deployments should be associated with
+	// This will be used by the agent to determine which Akita service to report traffic to
+	projectNameFlag string
+	// Represents the options for generating a secret
+	// When set to "false" or left empty, injectCmd will not generate a secret
+	// When set to "true", injectCmd will prepend a secret to each injectable namespace found in the file to inject (injectFileNameFlag)
+	// Otherwise, injectCmd will treat secretInjectFlag as the file path all secrets should be generated to
+	secretInjectFlag string
+)
+
+var injectCmd = &cobra.Command{
+	Use:   "inject",
+	Short: "Inject Akita into a Kubernetes deployment",
+	Long:  "Inject Akita into a Kubernetes deployment or set of deployments, and output the result to stdout or a file",
+	RunE: func(_ *cobra.Command, args []string) error {
+		secretOpts := resolveSecretGenerationOptions(secretInjectFlag)
+
+		// To avoid users unintentionally attempting to apply injected Deployments via pipeline without
+		// their dependent Secrets, require that the user explicitly specify an output file.
+		if secretOpts.ShouldInject && secretOpts.Filepath.IsSome() && injectOutputFlag == "" {
+			printer.Errorln("Cannot specify a Secret file path without an output file (using --output or -o)")
+			printer.Infoln("To generate a Secret file on its own, use `akita kube secret`")
+			return cmderr.AkitaErr{
+				Err: errors.New("invalid flag usage"),
+			}
+		}
+
+		// Create the injector which reads from the Kubernetes YAML file specified by the user
+		injectr, err := injector.FromYAML(injectFileNameFlag)
+		if err != nil {
+			return cmderr.AkitaErr{
+				Err: errors.Wrapf(
+					err,
+					"Failed to read injection file %s",
+					injectFileNameFlag,
+				),
+			}
+		}
+
+		// Generate a secret for each namespace in the deployment if the user specified secret generation
+		secretBuf := new(bytes.Buffer)
+		if secretOpts.ShouldInject {
+			key, secret, err := cmderr.RequireAPICredentials("API credentials are required to generate secret.")
+			if err != nil {
+				return err
+			}
+
+			namespaces, err := injectr.InjectableNamespaces()
+			if err != nil {
+				return err
+			}
+
+			for _, namespace := range namespaces {
+				r, err := handleSecretGeneration(namespace, key, secret)
+				if err != nil {
+					return err
+				}
+
+				secretBuf.WriteString("---\n")
+				secretBuf.Write(r)
+			}
+		}
+
+		// Create the output buffer
+		out := new(bytes.Buffer)
+
+		// Either write the secret to a file or prepend it to the output
+		if secretFilePath, exists := secretOpts.Filepath.Get(); exists {
+			err = writeFile(secretBuf.Bytes(), secretFilePath)
+			if err != nil {
+				return err
+			}
+
+			printer.Infof("Kubernetes Secret generated to %s\n", secretFilePath)
+		} else {
+			// Assign the secret to the output buffer
+			// We do this so that the secret is written before any injected Deployment resources that depend on it
+			out = secretBuf
+		}
+
+		// Inject the sidecar into the input file
+		rawInjected, err := injector.ToRawYAML(injectr, createSidecar(projectNameFlag))
+		if err != nil {
+			return cmderr.AkitaErr{Err: errors.Wrap(err, "Failed to inject sidecars")}
+		}
+		// Append the injected YAML to the output
+		out.Write(rawInjected)
+
+		// If the user did not specify an output file, print the output to stdout
+		if injectOutputFlag == "" {
+			printer.Stdout.RawOutput(out.String())
+			return nil
+		}
+
+		// Write the output to the specified file
+		if err := writeFile(out.Bytes(), injectOutputFlag); err != nil {
+			return err
+		}
+		printer.Infof("Injected YAML written to %s\n", injectOutputFlag)
+
+		return nil
+	},
+	PersistentPreRun: func(cmd *cobra.Command, args []string) {
+		// Initialize the telemetry client, but do not allow any logs to be printed
+		telemetry.Init(false)
+	},
+}
+
+// A parsed representation of the `--secret` option.
+type secretGenerationOptions struct {
+	// Whether to inject a secret
+	ShouldInject bool
+	// The path to the secret file
+	Filepath optionals.Optional[string]
+}
+
+func createSidecar(projectName string) v1.Container {
+	sidecar := v1.Container{
+		Name:  "akita",
+		Image: "akitasoftware/cli:latest",
+		Env: []v1.EnvVar{
+			{
+				Name: "AKITA_API_KEY_ID",
+				ValueFrom: &v1.EnvVarSource{
+					SecretKeyRef: &v1.SecretKeySelector{
+						LocalObjectReference: v1.LocalObjectReference{
+							Name: "akita-secrets",
+						},
+						Key: "akita-api-key",
+					},
+				},
+			},
+			{
+				Name: "AKITA_API_KEY_SECRET",
+				ValueFrom: &v1.EnvVarSource{
+					SecretKeyRef: &v1.SecretKeySelector{
+						LocalObjectReference: v1.LocalObjectReference{
+							Name: "akita-secrets",
+						},
+						Key: "akita-api-secret",
+					},
+				},
+			},
+		},
+		Lifecycle: &v1.Lifecycle{
+			PreStop: &v1.LifecycleHandler{
+				Exec: &v1.ExecAction{
+					Command: []string{
+						"/bin/sh",
+						"-c",
+						"AKITA_PID=$(pgrep akita) && kill -2 $AKITA_PID && tail -f /proc/$AKITA_PID/fd/1",
+					},
+				},
+			},
+		},
+		Args: []string{"apidump", "--project", projectName},
+		SecurityContext: &v1.SecurityContext{
+			Capabilities: &v1.Capabilities{Add: []v1.Capability{"NET_RAW"}},
+		},
+	}
+
+	return sidecar
+}
+
+// Parses the given value for the `--secret` option.
+func resolveSecretGenerationOptions(flagValue string) secretGenerationOptions {
+	if flagValue == "" || flagValue == "false" {
+		return secretGenerationOptions{
+			ShouldInject: false,
+			Filepath:     optionals.None[string](),
+		}
+	}
+
+	if flagValue == "true" {
+		return secretGenerationOptions{
+			ShouldInject: true,
+			Filepath:     optionals.None[string](),
+		}
+	}
+
+	return secretGenerationOptions{
+		ShouldInject: true,
+		Filepath:     optionals.Some(flagValue),
+	}
+}
+
+func init() {
+	injectCmd.Flags().StringVarP(
+		&injectFileNameFlag,
+		"file",
+		"f",
+		"",
+		"Path to the Kubernetes YAML file to be injected. This should contain a Deployment object.",
+	)
+	_ = injectCmd.MarkFlagRequired("file")
+
+	injectCmd.Flags().StringVarP(
+		&injectOutputFlag,
+		"output",
+		"o",
+		"",
+		"Path to the output file. If not specified, the output will be printed to stdout.",
+	)
+
+	injectCmd.Flags().StringVarP(
+		&projectNameFlag,
+		"project",
+		"p",
+		"",
+		"Name of the Akita project to which the traffic will be uploaded.",
+	)
+	_ = injectCmd.MarkFlagRequired("project")
+
+	injectCmd.Flags().StringVarP(
+		&secretInjectFlag,
+		"secret",
+		"s",
+		"false",
+		`Whether to generate a Kubernetes Secret. If set to "true", the secret will be added to the modified Kubernetes YAML file. Specify a path to write the secret to a separate file; if this is done, an output file must also be specified with --output.`,
+	)
+	// Default value is "true" when the flag is given without an argument.
+	injectCmd.Flags().Lookup("secret").NoOptDefVal = "true"
+
+	Cmd.AddCommand(injectCmd)
+}

cmd/internal/kube/injector/helper.go

@@ -0,0 +1,28 @@
+package injector
+
+import (
+	"bytes"
+	v1 "k8s.io/api/core/v1"
+	kyaml "sigs.k8s.io/yaml"
+)
+
+// Calls the given injector's Inject method and returns the result as a YAML bytes.
+func ToRawYAML(injector Injector, sidecar v1.Container) ([]byte, error) {
+	injectedObjects, err := injector.Inject(sidecar)
+	if err != nil {
+		return nil, err
+	}
+
+	out := new(bytes.Buffer)
+	for _, obj := range injectedObjects {
+		raw, err := kyaml.Marshal(obj)
+		if err != nil {
+			return nil, err
+		}
+
+		out.WriteString("---\n")
+		out.Write(raw)
+	}
+
+	return out.Bytes(), nil
+}