sshpoke/cmd/root.go

package cmd

import (
	"context"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/Neur0toxine/sshpoke/internal/config"
	"github.com/Neur0toxine/sshpoke/internal/docker"
	"github.com/Neur0toxine/sshpoke/internal/logger"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

var cfgFile string

var rootCmd = &cobra.Command{
	Use:   "sshpoke",
	Short: "Expose your Docker services to the Internet via SSH.",
	Long:  `sshpoke is a CLI application that listens to the docker socket and automatically exposes relevant services to the Internet.`,
	Run: func(cmd *cobra.Command, args []string) {
		var err error
		ctx, cancel := context.WithCancel(context.Background())
		docker.Default, err = docker.New(ctx)
		if err != nil {
			logger.Sugar.Fatalf("cannot connect to docker daemon: %s", err)
		}

		for id, item := range docker.Default.Containers() {
			logger.Sugar.Debugw("registering container",
				"id", id,
				"ip", item.IP.String(),
				"port", item.Port,
				"server", item.Server,
				"domain", item.Domain)
		}

		events, err := docker.Default.Listen()
		if err != nil {
			logger.Sugar.Fatalf("cannot listen to docker events: %s", err)
		}

		go func() {
			logger.Sugar.Debug("listening for docker events...")
			for event := range events {
				_ = event
			}
		}()

		linuxSig := make(chan os.Signal, 1)
		signal.Notify(linuxSig)
		for sig := range linuxSig {
			switch sig {
			case os.Interrupt, syscall.SIGQUIT, syscall.SIGTERM:
				cancel()
				time.Sleep(time.Millisecond * 200)
				logger.Sugar.Infof("received %s, exiting...", sig)
				os.Exit(0)
			default:
			}
		}
	},
}

func Execute() {
	err := rootCmd.Execute()
	if err != nil {
		os.Exit(1)
	}
}

func init() {
	cobra.OnInitialize(initConfig)
	rootCmd.PersistentFlags().StringVar(&cfgFile, "config", "config.yml", "Configuration file (default is config.yml)")
}

func initConfig() {
	if cfgFile != "" {
		viper.SetConfigFile(cfgFile)
	} else {
		workingDir, err := os.Getwd()
		cobra.CheckErr(err)
		viper.AddConfigPath(workingDir)
		viper.SetConfigType("yml")
		viper.SetConfigName("config")
	}

	log := logger.New(os.Getenv("SSHPOKE_DEBUG") == "true").Sugar()
	viper.SetEnvPrefix("SSHPOKE")
	viper.AutomaticEnv()
	if err := config.BindStructEnv(&config.DefaultConfig); err != nil {
		log.Fatalf("cannot bind configuration keys: %s", err)
	}
	if err := viper.ReadInConfig(); err == nil {
		log.Debugf("using config file: %s", viper.ConfigFileUsed())
	}
	if err := viper.Unmarshal(&config.DefaultConfig); err != nil {
		log.Fatalf("cannot load configuration: %s", err)
	}
	logger.Initialize()
}