apiary/cmd/apiary.go

package main

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

	"github.com/urfave/cli/v2"
	"github.uio.no/torjus/apiary"
	"github.uio.no/torjus/apiary/config"
	"github.uio.no/torjus/apiary/honeypot"
	"github.uio.no/torjus/apiary/honeypot/store"
	"github.uio.no/torjus/apiary/web"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
)

func main() {
	app := &cli.App{
		Name:    "apiary",
		Version: apiary.Version,
		Authors: []*cli.Author{
			{
				Name:  "Torjus Håkestad",
				Email: "torjus@usit.uio.no",
			},
		},
		Commands: []*cli.Command{
			{
				Name:   "serve",
				Action: ActionServe,
				Usage:  "Start Apiary server",
			},
		},
	}

	if err := app.Run(os.Args); err != nil {
		fmt.Printf("Error: %s\n", err)
		os.Exit(1)
	}
}

func ActionServe(c *cli.Context) error {
	cfg, err := getConfig()
	if err != nil {
		return err
	}

	loggers := setupLoggers(cfg)

	var s store.LoginAttemptStore
	switch cfg.Store.Type {
	case "MEMORY", "memory":
		s = &store.MemoryStore{}
	case "POSTGRES", "postgres":
		pgStore, err := store.NewPostgresStore(cfg.Store.Postgres.DSN)
		if err != nil {
			return err
		}
		if err := pgStore.InitDB(); err != nil {
			return err
		}
		s = pgStore
	default:
		return fmt.Errorf("Invalid store configured")
	}

	hs, err := honeypot.NewHoneypotServer(cfg.Honeypot, s)
	if err != nil {
		return err
	}
	hs.Logger = loggers.honeypotLogger

	web := web.NewServer(cfg.Frontend, hs, s)
	web.AccessLogger = loggers.webAccessLogger
	web.ServerLogger = loggers.webServerLogger

	interruptChan := make(chan os.Signal, 1)
	signal.Notify(interruptChan, os.Interrupt)

	rootCtx, rootCancel := context.WithCancel(c.Context)
	serversCtx, serversCancel := context.WithCancel(rootCtx)

	// Handle interrupt
	go func() {
		<-interruptChan
		fmt.Println("Got interrupt. Shutting down.")
		serversCancel()
	}()

	// Start ssh server
	go func() {
		hs.ListenAndServe()
	}()

	// Start web server
	go func() {
		web.ListenAndServe()
	}()

	go func() {
		<-serversCtx.Done()

		// Stop SSH server
		sshShutdownCtx, sshShutdownCancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer sshShutdownCancel()
		loggers.rootLogger.Info("SSH server shutdown started")
		if err := hs.Shutdown(sshShutdownCtx); err != nil {
			loggers.rootLogger.Infow("Error shutting down SSH server", "error", err)
		}
		loggers.rootLogger.Info("SSH server shutdown complete")

		// Stop Web server
		webShutdownCtx, webShutdownCancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer webShutdownCancel()

		loggers.rootLogger.Info("Web server shutdown started")
		if err := web.Shutdown(webShutdownCtx); err != nil {
			loggers.rootLogger.Infow("Error shutting down web server", "error", err)
		}
		loggers.rootLogger.Info("Web server shutdown complete")
		rootCancel()
	}()

	<-rootCtx.Done()

	return nil

}

type loggerCollection struct {
	rootLogger      *zap.SugaredLogger
	honeypotLogger  *zap.SugaredLogger
	webAccessLogger *zap.SugaredLogger
	webServerLogger *zap.SugaredLogger
}

func setupLoggers(cfg config.Config) *loggerCollection {
	logEncoderCfg := zap.NewProductionEncoderConfig()
	logEncoderCfg.EncodeCaller = func(caller zapcore.EntryCaller, enc zapcore.PrimitiveArrayEncoder) {}

	level := zap.NewAtomicLevelAt(zap.InfoLevel)
	logEncoderCfg.EncodeLevel = zapcore.CapitalColorLevelEncoder
	logEncoderCfg.EncodeTime = zapcore.ISO8601TimeEncoder
	logEncoderCfg.EncodeDuration = zapcore.NanosDurationEncoder
	rootLoggerCfg := &zap.Config{
		Level:            level,
		OutputPaths:      []string{"stdout"},
		ErrorOutputPaths: []string{"stderr"},
		Encoding:         "console",
		EncoderConfig:    logEncoderCfg,
	}
	rootLogger, err := rootLoggerCfg.Build()
	if err != nil {
		panic(err)
	}

	return &loggerCollection{
		rootLogger:      rootLogger.Named("APP").Sugar(),
		honeypotLogger:  rootLogger.Named("HON").Sugar(),
		webAccessLogger: rootLogger.Named("ACC").Sugar(),
		webServerLogger: rootLogger.Named("WEB").Sugar(),
	}

}

func getConfig() (config.Config, error) {
	defaultLocations := []string{
		"apiary.toml",
		"/etc/apiary.toml",
	}

	for _, fname := range defaultLocations {
		if _, err := os.Stat(fname); os.IsNotExist(err) {
			continue
		}
		cfg, err := config.FromFile(fname)
		if err != nil {
			return config.Config{}, err
		}
		return cfg, nil
	}
	return config.Config{}, fmt.Errorf("Could not find config file")
}