Initial commit

cmd/apiary.go

@@ -0,0 +1,186 @@
+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")
+}