cli/packages/cmd/gateway.go

package cmd

import (
	"context"
	"fmt"
	"os"
	"os/exec"
	"os/signal"
	"runtime"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/hanzokms/cli/packages/api"
	"github.com/hanzokms/cli/packages/config"
	"github.com/hanzokms/cli/packages/gateway"
	gatewayv2 "github.com/hanzokms/cli/packages/gateway-v2"
	"github.com/hanzokms/cli/packages/pam/session"
	"github.com/hanzokms/cli/packages/util"
	kmsSdk "github.com/hanzokms/go-sdk"
	"github.com/pkg/errors"
	insights "github.com/hanzoai/insights-go"
	"github.com/rs/zerolog/log"
	"github.com/spf13/cobra"
)

func getKMSSdkInstance(cmd *cobra.Command) (kmsSdk.ClientInterface, context.CancelFunc, error) {

	ctx, cancel := context.WithCancel(cmd.Context())
	kmsClient := kmsSdk.NewClient(ctx, kmsSdk.Config{
		SiteUrl:   config.KMS_URL,
		UserAgent: api.USER_AGENT,
	})

	token, err := util.GetKMSToken(cmd)
	if err != nil {
		cancel()
		return nil, nil, err
	}

	// if the --token param is set, we use it directly for authentication
	if token != nil {
		kmsClient.Auth().SetAccessToken(token.Token)
		return kmsClient, cancel, nil
	}

	// if the --token param is not set, we use the auth-method flag to determine the authentication method, and perform the appropriate login flow based on that
	authMethod, err := util.GetCmdFlagOrEnv(cmd, "auth-method", []string{util.KMS_AUTH_METHOD_NAME})

	if err != nil {
		cancel()
		return nil, nil, err
	}

	authMethodValid, strategy := util.IsAuthMethodValid(authMethod, false)
	if !authMethodValid {
		util.PrintErrorMessageAndExit(fmt.Sprintf("Invalid login method: %s", authMethod))
	}

	sdkAuthenticator := util.NewSdkAuthenticator(kmsClient, cmd)

	authStrategies := map[util.AuthStrategyType]func() (credential kmsSdk.MachineIdentityCredential, e error){
		util.AuthStrategy.UNIVERSAL_AUTH:    sdkAuthenticator.HandleUniversalAuthLogin,
		util.AuthStrategy.KUBERNETES_AUTH:   sdkAuthenticator.HandleKubernetesAuthLogin,
		util.AuthStrategy.AZURE_AUTH:        sdkAuthenticator.HandleAzureAuthLogin,
		util.AuthStrategy.GCP_ID_TOKEN_AUTH: sdkAuthenticator.HandleGcpIdTokenAuthLogin,
		util.AuthStrategy.GCP_IAM_AUTH:      sdkAuthenticator.HandleGcpIamAuthLogin,
		util.AuthStrategy.AWS_IAM_AUTH:      sdkAuthenticator.HandleAwsIamAuthLogin,
		util.AuthStrategy.OIDC_AUTH:         sdkAuthenticator.HandleOidcAuthLogin,
		util.AuthStrategy.JWT_AUTH:          sdkAuthenticator.HandleJwtAuthLogin,
	}

	_, err = authStrategies[strategy]()

	if err != nil {
		cancel()
		return nil, nil, err
	}

	return kmsClient, cancel, nil
}

var gatewayCmd = &cobra.Command{
	Use:   "gateway",
	Short: "Run the KMS gateway or manage its systemd service",
	Long:  "Run the KMS gateway in the foreground or manage its systemd service installation. Use 'gateway install' to set up the systemd service.",
	Example: `kms gateway --token=<token>
  sudo kms gateway install --token=<token> --domain=<domain>`,
	DisableFlagsInUseLine: true,
	Args:                  cobra.NoArgs,
	Run: func(cmd *cobra.Command, args []string) {
		log.Info().Msg("DEPRECATION NOTICE: The 'kms gateway' command will be deprecated in a future version. Please use 'kms gateway start'.\nNOTE: This requires manually updating your existing resources to point to the new gateway.")

		kmsClient, cancelSdk, err := getKMSSdkInstance(cmd)
		if err != nil {
			util.HandleError(err, "unable to get KMS client")
		}
		defer cancelSdk()

		var accessToken atomic.Value
		accessToken.Store(kmsClient.Auth().GetAccessToken())

		if accessToken.Load().(string) == "" {
			util.HandleError(errors.New("no access token found"))
		}

		Telemetry.CaptureEvent("cli-command:gateway", insights.NewProperties().Set("version", util.CLI_VERSION))

		sigCh := make(chan os.Signal, 1)
		signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
		sigStopCh := make(chan bool, 1)

		ctx, cancelCmd := context.WithCancel(cmd.Context())
		defer cancelCmd()

		go func() {
			<-sigCh
			close(sigStopCh)
			cancelCmd()
			cancelSdk()

			// If we get a second signal, force exit
			<-sigCh
			log.Warn().Msgf("Force exit triggered")
			os.Exit(1)
		}()

		var gatewayInstance *gateway.Gateway

		// Token refresh goroutine - runs every 10 seconds
		go func() {
			tokenRefreshTicker := time.NewTicker(10 * time.Second)
			defer tokenRefreshTicker.Stop()

			for {
				select {
				case <-tokenRefreshTicker.C:
					if ctx.Err() != nil {
						return
					}

					newToken := kmsClient.Auth().GetAccessToken()
					if newToken != "" && newToken != accessToken.Load().(string) {
						accessToken.Store(newToken)
						if gatewayInstance != nil {
							gatewayInstance.UpdateIdentityAccessToken(newToken)
						}
					}

				case <-ctx.Done():
					return
				}
			}
		}()

		// Main gateway retry loop with proper context handling
		retryTicker := time.NewTicker(5 * time.Second)
		defer retryTicker.Stop()

		for {
			if ctx.Err() != nil {
				log.Info().Msg("Shutting down gateway")
				return
			}
			gatewayInstance, err := gateway.NewGateway(accessToken.Load().(string))
			if err != nil {
				util.HandleError(err)
			}

			if err = gatewayInstance.ConnectWithRelay(); err != nil {
				if ctx.Err() != nil {
					log.Info().Msg("Shutting down gateway")
					return
				}

				log.Error().Msgf("Gateway connection error with relay: %s", err)
				log.Info().Msg("Retrying connection in 5 seconds...")
				select {
				case <-retryTicker.C:
					continue
				case <-ctx.Done():
					log.Info().Msg("Shutting down gateway")
					return
				}
			}

			err = gatewayInstance.Listen(ctx)
			if ctx.Err() != nil {
				log.Info().Msg("Gateway shutdown complete")
				return
			}
			log.Error().Msgf("Gateway listen error: %s", err)
			log.Info().Msg("Retrying connection in 5 seconds...")
			select {
			case <-retryTicker.C:
				continue
			case <-ctx.Done():
				log.Info().Msg("Shutting down gateway")
				return
			}
		}
	},
}

var gatewayStartCmd = &cobra.Command{
	Use:                   "start",
	Short:                 "Start the new KMS gateway",
	Long:                  "Start the new KMS gateway component.",
	Example:               "kms gateway start --name=<name> --token=<token>",
	DisableFlagsInUseLine: true,
	Args:                  cobra.NoArgs,
	Run: func(cmd *cobra.Command, args []string) {
		gatewayName, err := util.GetCmdFlagOrEnv(cmd, "name", []string{gatewayv2.GATEWAY_NAME_ENV_NAME})
		if err != nil {
			util.HandleError(err, fmt.Sprintf("unable to get name flag or %s env", gatewayv2.GATEWAY_NAME_ENV_NAME))
		}

		pamSessionRecordingPath, err := util.GetCmdFlagOrEnv(cmd, "pam-session-recording-path", []string{gatewayv2.KMS_PAM_SESSION_RECORDING_PATH_ENV_NAME})
		if err == nil && pamSessionRecordingPath != "" {
			session.SetSessionRecordingPath(pamSessionRecordingPath)
		}

		kmsClient, cancelSdk, err := getKMSSdkInstance(cmd)
		if err != nil {
			util.HandleError(err, "unable to get KMS client")
		}
		defer cancelSdk()

		var accessToken atomic.Value
		accessToken.Store(kmsClient.Auth().GetAccessToken())

		if accessToken.Load().(string) == "" {
			util.HandleError(errors.New("no access token found"))
		}

		relayName, err := util.GetRelayName(cmd, false, accessToken.Load().(string))
		if err != nil {
			util.HandleError(err, "unable to get relay name")
		}

		gatewayInstance, err := gatewayv2.NewGateway(&gatewayv2.GatewayConfig{
			Name:           gatewayName,
			RelayName:      relayName,
			ReconnectDelay: 10 * time.Second,
		})

		if err != nil {
			util.HandleError(err, "unable to create gateway instance")
		}

		gatewayInstance.SetToken(accessToken.Load().(string))

		Telemetry.CaptureEvent("cli-command:gateway-v2", insights.NewProperties().Set("version", util.CLI_VERSION))

		sigCh := make(chan os.Signal, 1)
		signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)

		ctx, cancelCmd := context.WithCancel(cmd.Context())
		defer cancelCmd()

		go func() {
			<-sigCh
			log.Info().Msg("Received shutdown signal, shutting down gateway...")
			cancelCmd()
			cancelSdk()

			// Give graceful shutdown 10 seconds, then force exit on second signal
			select {
			case <-sigCh:
				log.Warn().Msg("Second signal received, force exit triggered")
				os.Exit(1)
			case <-time.After(10 * time.Second):
				log.Info().Msg("Graceful shutdown completed")
				os.Exit(0)
			}
		}()

		// Token refresh goroutine - runs every 10 seconds
		go func() {
			tokenRefreshTicker := time.NewTicker(10 * time.Second)
			defer tokenRefreshTicker.Stop()

			for {
				select {
				case <-tokenRefreshTicker.C:
					if ctx.Err() != nil {
						return
					}

					newToken := kmsClient.Auth().GetAccessToken()
					if newToken != "" && newToken != accessToken.Load().(string) {
						accessToken.Store(newToken)
						gatewayInstance.SetToken(newToken)
					}

				case <-ctx.Done():
					return
				}
			}
		}()

		err = gatewayInstance.Start(ctx)
		if err != nil {
			util.HandleError(err, "unable to start gateway instance")
		}
	},
}

var gatewayInstallCmd = &cobra.Command{
	Use:                   "install",
	Short:                 "Install and enable systemd service for the gateway (requires sudo)",
	Long:                  "Install and enable systemd service for the gateway. Must be run with sudo on Linux.",
	Example:               "sudo kms gateway install --token=<token> --domain=<domain>",
	DisableFlagsInUseLine: true,
	Args:                  cobra.NoArgs,
	Run: func(cmd *cobra.Command, args []string) {
		if runtime.GOOS != "linux" {
			util.HandleError(fmt.Errorf("systemd service installation is only supported on Linux"))
		}

		if os.Geteuid() != 0 {
			util.HandleError(fmt.Errorf("systemd service installation requires root/sudo privileges"))
		}

		token, err := util.GetKMSToken(cmd)
		if err != nil {
			util.HandleError(err, "Unable to parse flag")
		}

		if token == nil {
			util.HandleError(errors.New("Token not found"))
		}

		domain, err := cmd.Flags().GetString("domain")
		if err != nil {
			util.HandleError(err, "Unable to parse domain flag")
		}

		if err := gateway.InstallGatewaySystemdService(token.Token, domain); err != nil {
			util.HandleError(err, "Failed to install systemd service")
		}

		enableCmd := exec.Command("systemctl", "enable", "kms-gateway")
		if err := enableCmd.Run(); err != nil {
			util.HandleError(err, "Failed to enable systemd service")
		}

		log.Info().Msg("Successfully installed and enabled kms-gateway service")
		log.Info().Msg("To start the service, run: sudo systemctl start kms-gateway")
	},
}

var gatewayUninstallCmd = &cobra.Command{
	Use:                   "uninstall",
	Short:                 "Uninstall and remove systemd service for the gateway (requires sudo)",
	Long:                  "Uninstall and remove systemd service for the gateway. Must be run with sudo on Linux.",
	Example:               "sudo kms gateway uninstall",
	DisableFlagsInUseLine: true,
	Args:                  cobra.NoArgs,
	Run: func(cmd *cobra.Command, args []string) {
		if runtime.GOOS != "linux" {
			util.HandleError(fmt.Errorf("systemd service installation is only supported on Linux"))
		}

		if os.Geteuid() != 0 {
			util.HandleError(fmt.Errorf("systemd service installation requires root/sudo privileges"))
		}

		if err := gateway.UninstallGatewaySystemdService(); err != nil {
			util.HandleError(err, "Failed to uninstall systemd service")
		}
	},
}

var gatewaySystemdCmd = &cobra.Command{
	Use:   "systemd",
	Short: "Manage systemd service for KMS gateway",
	Long:  "Manage systemd service for KMS gateway. Use 'systemd install' to install and enable the service.",
	Example: `sudo kms gateway systemd install --token=<token> --domain=<domain> --name=<name>
  sudo kms gateway systemd uninstall`,
	DisableFlagsInUseLine: true,
	Args:                  cobra.NoArgs,
}

var gatewaySystemdInstallCmd = &cobra.Command{
	Use:                   "install",
	Short:                 "Install and enable systemd service for the gateway (v2) (requires sudo)",
	Long:                  "Install and enable systemd service for the new gateway (v2). Must be run with sudo on Linux.",
	Example:               "sudo kms gateway systemd install --token=<token> --domain=<domain> --name=<name>",
	DisableFlagsInUseLine: true,
	Args:                  cobra.NoArgs,
	Run: func(cmd *cobra.Command, args []string) {
		if runtime.GOOS != "linux" {
			util.HandleError(fmt.Errorf("systemd service installation is only supported on Linux"))
		}

		if os.Geteuid() != 0 {
			util.HandleError(fmt.Errorf("systemd service installation requires root/sudo privileges"))
		}

		token, err := util.GetKMSToken(cmd)
		if err != nil {
			util.HandleError(err, "Unable to parse flag")
		}

		if token == nil {
			util.HandleError(errors.New("Token not found"))
		}

		domain, err := cmd.Flags().GetString("domain")
		if err != nil {
			util.HandleError(err, "Unable to parse domain flag")
		}

		if domain != "" {
			config.KMS_URL = util.AppendAPIEndpoint(domain)
		}

		gatewayName, err := cmd.Flags().GetString("name")
		if err != nil {
			util.HandleError(err, "Unable to parse name flag")
		}
		if gatewayName == "" {
			util.HandleError(errors.New("Gateway name is required"))
		}

		serviceLogFile, err := cmd.Flags().GetString("log-file")
		if err != nil {
			util.HandleError(err, "Unable to parse log-file flag")
		}

		relayName, err := util.GetRelayName(cmd, false, token.Token)
		if err != nil {
			util.HandleError(err, "unable to get relay name")
		}

		err = gatewayv2.InstallGatewaySystemdService(token.Token, domain, gatewayName, relayName, serviceLogFile)
		if err != nil {
			util.HandleError(err, "Unable to install systemd service")
		}

		enableCmd := exec.Command("systemctl", "enable", "kms-gateway")
		if err := enableCmd.Run(); err != nil {
			util.HandleError(err, "Failed to enable systemd service")
		}

		log.Info().Msg("Successfully installed and enabled kms-gateway service")
		log.Info().Msg("To start the service, run: sudo systemctl start kms-gateway")
	},
}

var gatewaySystemdUninstallCmd = &cobra.Command{
	Use:                   "uninstall",
	Short:                 "Uninstall and remove systemd service for the gateway (requires sudo)",
	Long:                  "Uninstall and remove systemd service for the gateway. Must be run with sudo on Linux.",
	Example:               "sudo kms gateway systemd uninstall",
	DisableFlagsInUseLine: true,
	Args:                  cobra.NoArgs,
	Run: func(cmd *cobra.Command, args []string) {
		if runtime.GOOS != "linux" {
			util.HandleError(fmt.Errorf("systemd service installation is only supported on Linux"))
		}

		if os.Geteuid() != 0 {
			util.HandleError(fmt.Errorf("systemd service installation requires root/sudo privileges"))
		}

		if err := gatewayv2.UninstallGatewaySystemdService(); err != nil {
			util.HandleError(err, "Failed to uninstall systemd service")
		}
	},
}

var gatewayRelayCmd = &cobra.Command{
	Example:               `kms gateway relay`,
	Short:                 "Used to run kms gateway relay",
	Use:                   "relay",
	DisableFlagsInUseLine: true,
	Args:                  cobra.NoArgs,
	Run: func(cmd *cobra.Command, args []string) {
		relayConfigFilePath, err := cmd.Flags().GetString("config")
		if err != nil {
			util.HandleError(err, "Unable to parse flag")
		}

		if relayConfigFilePath == "" {
			util.HandleError(errors.New("Missing config file"))
		}

		gatewayRelay, err := gateway.NewGatewayRelay(relayConfigFilePath)
		if err != nil {
			util.HandleError(err, "Failed to initialize gateway")
		}
		err = gatewayRelay.Run()
		if err != nil {
			util.HandleError(err, "Failed to start gateway")
		}
	},
}

func init() {
	// Legacy gateway command flags (v1)
	gatewayCmd.Flags().String("token", "", "connect with Hanzo KMS using machine identity access token. if not provided, you must set the auth-method flag")
	gatewayCmd.Flags().String("auth-method", "", "login method [universal-auth, kubernetes, azure, gcp-id-token, gcp-iam, aws-iam, oidc-auth]. if not provided, you must set the token flag")
	gatewayCmd.Flags().String("client-id", "", "client id for universal auth")
	gatewayCmd.Flags().String("client-secret", "", "client secret for universal auth")
	gatewayCmd.Flags().String("machine-identity-id", "", "machine identity id for kubernetes, azure, gcp-id-token, gcp-iam, and aws-iam auth methods")
	gatewayCmd.Flags().String("service-account-token-path", "", "service account token path for kubernetes auth")
	gatewayCmd.Flags().String("service-account-key-file-path", "", "service account key file path for GCP IAM auth")
	gatewayCmd.Flags().String("jwt", "", "JWT for jwt-based auth methods [oidc-auth, jwt-auth]")

	// Gateway start command flags (v2)
	gatewayStartCmd.Flags().String("relay", "", "name of the relay to connect to (deprecated, use --target-relay-name)") // Deprecated, use --target-relay-name instead
	gatewayStartCmd.Flags().String("target-relay-name", "", "name of the relay to connect to")
	gatewayStartCmd.Flags().String("name", "", "name of the gateway")
	gatewayStartCmd.Flags().String("token", "", "connect with Hanzo KMS using machine identity access token. if not provided, you must set the auth-method flag")
	gatewayStartCmd.Flags().String("auth-method", "", "login method [universal-auth, kubernetes, azure, gcp-id-token, gcp-iam, aws-iam, oidc-auth]. if not provided, you must set the token flag")
	gatewayStartCmd.Flags().String("organization-slug", "", "When set, this will scope the login session to the specified sub-organization the machine identity has access to. If left empty, the session defaults to the organization where the machine identity was created in.")
	gatewayStartCmd.Flags().String("client-id", "", "client id for universal auth")
	gatewayStartCmd.Flags().String("client-secret", "", "client secret for universal auth")
	gatewayStartCmd.Flags().String("machine-identity-id", "", "machine identity id for kubernetes, azure, gcp-id-token, gcp-iam, and aws-iam auth methods")
	gatewayStartCmd.Flags().String("service-account-token-path", "", "service account token path for kubernetes auth")
	gatewayStartCmd.Flags().String("service-account-key-file-path", "", "service account key file path for GCP IAM auth")
	gatewayStartCmd.Flags().String("jwt", "", "JWT for jwt-based auth methods [oidc-auth, jwt-auth]")
	gatewayStartCmd.Flags().String("pam-session-recording-path", "", "directory path for PAM session recordings (defaults to /var/lib/hanzo-kms/session_recordings)")

	// Legacy install command flags (v1)
	gatewayInstallCmd.Flags().String("token", "", "Connect with Hanzo KMS using machine identity access token")
	gatewayInstallCmd.Flags().String("domain", "", "Domain of your self-hosted Hanzo KMS instance")

	// Systemd install command flags (v2)
	gatewaySystemdInstallCmd.Flags().String("token", "", "Connect with Hanzo KMS using machine identity access token")
	gatewaySystemdInstallCmd.Flags().String("domain", "", "Domain of your self-hosted Hanzo KMS instance")
	gatewaySystemdInstallCmd.Flags().String("name", "", "The name of the gateway")
	gatewaySystemdInstallCmd.Flags().String("relay", "", "The name of the relay (deprecated, use --target-relay-name)") // Deprecated, use --target-relay-name instead
	gatewaySystemdInstallCmd.Flags().String("target-relay-name", "", "The name of the relay")
	gatewaySystemdInstallCmd.Flags().String("log-file", "", "The file to write the service logs to. Example: /var/log/hanzo-kms/gateway.log. If not provided, logs will not be written to a file.")

	// Gateway relay command flags
	gatewayRelayCmd.Flags().String("config", "", "Relay config yaml file path")

	// Wire up command hierarchy
	gatewaySystemdCmd.AddCommand(gatewaySystemdInstallCmd)
	gatewaySystemdCmd.AddCommand(gatewaySystemdUninstallCmd)

	gatewayCmd.AddCommand(gatewayStartCmd)
	gatewayCmd.AddCommand(gatewaySystemdCmd)
	gatewayCmd.AddCommand(gatewayInstallCmd)
	gatewayCmd.AddCommand(gatewayUninstallCmd)
	gatewayCmd.AddCommand(gatewayRelayCmd)
	RootCmd.AddCommand(gatewayCmd)
}