sshpoke/cmd/root.go

package cmd

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

	"github.com/Neur0toxine/sshpoke/internal/api/plugin"
	"github.com/Neur0toxine/sshpoke/internal/api/rest"
	"github.com/Neur0toxine/sshpoke/internal/config"
	"github.com/Neur0toxine/sshpoke/internal/docker"
	"github.com/Neur0toxine/sshpoke/internal/logger"
	"github.com/Neur0toxine/sshpoke/internal/server"
	"github.com/Neur0toxine/sshpoke/pkg/dto"
	plugin2 "github.com/Neur0toxine/sshpoke/pkg/plugin"
	"github.com/go-playground/validator/v10"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
)

var (
	cfgFile string
	Version string
)

var rootCmd = &cobra.Command{
	Use:     "sshpoke",
	Version: Version,
	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) {
		ctx, cancel := context.WithCancel(context.Background())
		server.DefaultManager = server.NewManager(ctx, config.Default.Servers, config.Default.DefaultServer)
		runPluginServer()
		runRestServer()
		runDockerEventListener(ctx)
		shutdown := makeShutdownFunc(cancel)

		linuxSig := make(chan os.Signal, 1)
		signal.Notify(linuxSig)
		for sig := range linuxSig {
			switch sig {
			case os.Interrupt, syscall.SIGQUIT, syscall.SIGTERM:
				shutdown(sig)
			default:
			}
		}
		shutdown(syscall.SIGHUP)
	},
}

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.Default); 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.Default); err != nil {
		log.Fatalf("cannot load configuration: %s", err)
	}
	if err := validator.New().Struct(config.Default); err != nil {
		log.Fatalf("invalid configuration: %s", err)
	}
	if config.Default.API.Rest.Port == 0 {
		config.Default.API.Rest.Port = rest.DefaultPort
	}
	if config.Default.API.Plugin.Port == 0 {
		config.Default.API.Plugin.Port = plugin2.DefaultPort
	}
	logger.Initialize()
	logger.Sugar.Debugw("configuration loaded", "config", config.Default)
}

func runPluginServer() {
	port := config.Default.API.Plugin.Port
	if port == 0 {
		port = plugin2.DefaultPort
	}
	go plugin.StartServer(port, logger.Sugar.With("component", "pluginServer"))
}

func runRestServer() {
	port := config.Default.API.Rest.Port
	if port == 0 {
		port = rest.DefaultPort
	}
	go rest.StartServer(
		config.Default.API.Rest.Token, port, logger.Sugar.With("component", "webServer"), config.Default.Debug)
}

func runDockerEventListener(ctx context.Context) {
	var err error
	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() {
		err := server.DefaultManager.ProcessEvent(dto.Event{
			Type:      dto.EventStart,
			Container: item,
		})
		if err != nil {
			logger.Sugar.Errorw("cannot expose container", "id", id, "error", err)
		}
	}

	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 {
			err := server.DefaultManager.ProcessEvent(event)
			if err != nil {
				logger.Sugar.Errorw("cannot expose container",
					"id", event.Container.ID, "error", err)
			}
		}
	}()
}

func makeShutdownFunc(cancel func()) func(os.Signal) {
	return func(sig os.Signal) {
		cancel()
		server.DefaultManager.WaitForShutdown()
		logger.Sugar.Infof("received %s, exiting...", sig)
		os.Exit(0)
	}
}