diff --git a/cmd/ecs_credential_provider.go b/cmd/ecs_credential_provider.go index da3db53..7b8c50b 100644 --- a/cmd/ecs_credential_provider.go +++ b/cmd/ecs_credential_provider.go @@ -38,7 +38,7 @@ func init() { var ecsCredentialProvider = &cobra.Command{ Use: "ecs_credential_provider", - Short: "Run a local ECS Credential Provider endpoint that serves and caches credentials for roles on demand", + Short: "RunService a local ECS Credential Provider endpoint that serves and caches credentials for roles on demand", RunE: runEcsMetadata, } @@ -65,8 +65,6 @@ func runEcsMetadata(cmd *cobra.Command, args []string) error { // Check for interrupt signal and exit cleanly <-shutdown log.Print("Shutdown signal received, exiting weep...") - // Send a signal to show that we're done shutting down - done <- 0 return nil } diff --git a/cmd/metadata.go b/cmd/metadata.go index 5be23e4..b23398b 100644 --- a/cmd/metadata.go +++ b/cmd/metadata.go @@ -42,7 +42,7 @@ func init() { var metadataCmd = &cobra.Command{ Use: "metadata [role_name]", - Short: "Run a local Instance Metadata Service (IMDS) endpoint that serves credentials", + Short: "RunService a local Instance Metadata Service (IMDS) endpoint that serves credentials", Args: cobra.ExactArgs(1), RunE: runMetadata, } @@ -86,8 +86,6 @@ func runMetadata(cmd *cobra.Command, args []string) error { // Check for interrupt signal and exit cleanly <-shutdown log.Print("Shutdown signal received, exiting weep meta-data service...") - // Send a signal to show that we're done shutting down - done <- 0 return nil } diff --git a/cmd/root.go b/cmd/root.go index 654fa58..2ea534c 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -19,9 +19,11 @@ package cmd import ( "fmt" "os" + "os/signal" "path" "runtime" "strings" + "syscall" "github.com/netflix/weep/util" @@ -41,8 +43,6 @@ var ( Long: "Weep is a CLI tool that manages AWS access via ConsoleMe for local development.", DisableAutoGenTag: true, } - shutdown chan os.Signal - done chan int ) func init() { @@ -53,15 +53,20 @@ func init() { rootCmd.PersistentFlags().StringSliceVarP(&assumeRole, "assume-role", "A", make([]string, 0), "one or more roles to assume after retrieving credentials") rootCmd.PersistentFlags().StringVar(&logFormat, "log-format", "", "log format (json or tty)") rootCmd.PersistentFlags().StringVar(&logLevel, "log-level", "", "log level (debug, info, warn)") - + rootCmd.PersistentFlags().BoolVarP(&runAsService, "svc", "s", false, "run weep as a service") } -func Execute(shutdownSignal chan os.Signal, doneSignal chan int) { - shutdown = shutdownSignal - done = doneSignal - if err := rootCmd.Execute(); err != nil { - _, _ = fmt.Fprintln(os.Stderr, err) - os.Exit(1) +func Execute() { + shutdown := make(chan os.Signal, 1) + done = make(chan int, 1) + signal.Notify(shutdown, syscall.SIGINT, syscall.SIGTERM) + if runAsService { + RunService() + } else { + if err := rootCmd.Execute(); err != nil { + _, _ = fmt.Fprintln(os.Stderr, err) + os.Exit(1) + } } } diff --git a/run/run.go b/cmd/svc.go similarity index 83% rename from run/run.go rename to cmd/svc.go index ab46a8f..ef780c6 100644 --- a/run/run.go +++ b/cmd/svc.go @@ -14,20 +14,14 @@ * limitations under the License. */ -package run +package cmd import ( - "os" - "os/signal" - "syscall" - "github.com/kardianos/service" - "github.com/netflix/weep/cmd" log "github.com/sirupsen/logrus" ) var svcLogger service.Logger -var done chan int type program struct{} @@ -37,10 +31,8 @@ func (p *program) Start(s service.Service) error { } func (p *program) run() { - shutdown := make(chan os.Signal, 1) - done = make(chan int, 1) - signal.Notify(shutdown, syscall.SIGINT, syscall.SIGTERM) - cmd.Execute(shutdown, done) + _ = rootCmd.Execute() + done <- 0 } func (p *program) Stop(s service.Service) error { @@ -48,7 +40,7 @@ func (p *program) Stop(s service.Service) error { return nil } -func Run() { +func RunService() { svcConfig := &service.Config{ Name: "Weep", DisplayName: "Weep", @@ -68,4 +60,5 @@ func Run() { if err != nil { _ = svcLogger.Error(err) } + done <- 0 } diff --git a/cmd/vars.go b/cmd/vars.go index 267b737..1e2cc89 100644 --- a/cmd/vars.go +++ b/cmd/vars.go @@ -16,6 +16,8 @@ package cmd +import "os" + var ( assumeRole []string role string @@ -32,4 +34,7 @@ var ( cfgFile string logLevel string logFormat string + runAsService bool + shutdown chan os.Signal + done chan int ) diff --git a/main.go b/main.go index 6654a17..69af277 100644 --- a/main.go +++ b/main.go @@ -19,7 +19,8 @@ package main import ( "os" - "github.com/netflix/weep/run" + "github.com/netflix/weep/cmd" + log "github.com/sirupsen/logrus" ) @@ -29,5 +30,5 @@ func init() { } func main() { - run.Run() + cmd.Execute() }