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") }