timetable-mgr/main.go

package main

import (
	"fmt"
	"os"
	"os/signal"
	"os/user"
	"syscall"
	"time"
	_ "time/tzdata"

	"git.fjla.uk/owlboard/timetable-mgr/background"
	"git.fjla.uk/owlboard/timetable-mgr/cif"
	"git.fjla.uk/owlboard/timetable-mgr/corpus"
	"git.fjla.uk/owlboard/timetable-mgr/dbAccess"
	"git.fjla.uk/owlboard/timetable-mgr/helpers"
	"git.fjla.uk/owlboard/timetable-mgr/log"
	"git.fjla.uk/owlboard/timetable-mgr/messaging"
	"git.fjla.uk/owlboard/timetable-mgr/pis"
	"git.fjla.uk/owlboard/timetable-mgr/stations"
	"git.fjla.uk/owlboard/timetable-mgr/vstp"
	"go.uber.org/zap"
)

const (
	bold  = "\033[1m"
	redB  = "\033[1;31m"
	blue  = "\033[32m" //Actually green!
	cyan  = "\033[36m"
	reset = "\033[0m"
)

func init() {
	printStartupBanner()
	fmt.Printf("%sVersion %s \n\n%s", bold+blue, helpers.Version, reset)

	// Exits is being run as root
	// not necessary and not secure
	checkRunAsRoot()
}

func main() {
	log.InitLogger()
	defer log.Cleanup()
	log.Info("Initialising OwlBoard timetable-mgr", zap.String("version", helpers.Version))
	cfg, err := helpers.LoadConfig()
	if err != nil {
		log.Fatal("Unable to load configuration", zap.Error(err))
		return
	}
	cfg.PrintConfig()

	dbAccess.InitDataAccess(cfg)
	dbAccess.PushVersionToDb()

	// Initialise a `stop` channel to signal goroutines to cleanup
	stop := make(chan struct{})

	// Start CIF Task Ticker
	background.InitTicker(cfg, stop)

	// Handle signals from the OS
	go handleSignals(cfg, stop)

	// Manually call CIF and CORPUS checks to ensure that they are
	// not delayed until the first ticker event.
	go cif.CheckCif(cfg)
	go corpus.CheckCorpus(cfg)
	go stations.Check()
	go pis.Check()

	if cfg.VstpOn {
		messaging.StompInit(cfg)
		vstp.Subscribe()
	}

	<-stop
}

// Traps SIGINT and SIGTERM signals and ensures cleanup() is run
func handleSignals(cfg *helpers.Configuration, stop chan<- struct{}) {
	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)

	sig := <-sigChan
	log.Warn("Signal received: " + sig.String())
	cleanup(cfg, stop)
}

// Cleans up open connections ready for a clean exit of the program
func cleanup(cfg *helpers.Configuration, stop chan<- struct{}) {
	log.Debug("Cleaning up open connections")
	if cfg.VstpOn {
		log.Info("Closing STOMP Client")
		messaging.Disconnect(messaging.Client)

	}
	if dbAccess.MongoClient != nil {
		log.Info("Closing MongoDB Client")
		dbAccess.CloseMongoClient()
	}
	log.Info("Signalling to other goroutines")
	close(stop)

	log.Info("Program ready to exit")

	time.Sleep(500 * time.Millisecond)

	log.Cleanup()

	os.Exit(0)
}

func printStartupBanner() {
	art := `
  ___           _ ____                      _                          
 / _ \__      _| | __ )  ___   __ _ _ __ __| |                         
| | | \ \ /\ / / |  _ \ / _ \ / _' | '__/ _' |                         
| |_| |\ V  V /| | |_) | (_) | (_| | | | (_| |                         
 \___/  \_/\_/ |_|____/ \___/ \__,_|_|  \__,_|                         

 _   _                _        _     _                                 
| |_(_)_ __ ___   ___| |_ __ _| |__ | | ___       _ __ ___   __ _ _ __ 
| __| | '_ ' _ \ / _ \ __/ _' | '_ \| |/ _ \_____| '_ ' _ \ / _' | '__|
| |_| | | | | | |  __/ || (_| | |_) | |  __/_____| | | | | | (_| | |   
 \__|_|_| |_| |_|\___|\__\__,_|_.__/|_|\___|     |_| |_| |_|\__, |_|   
                                                            |___/      
	`

	fmt.Println(cyan + art + reset)
}

func checkRunAsRoot() {
	uid := os.Getuid()
	if uid != 0 {
		currUser, err := user.Current()
		var msg string
		if err != nil {
			msg = "Unable to determine which user is running the application."
		} else {
			msg = fmt.Sprintf("Running as user: %s, %s", currUser.Uid, currUser.Username)
		}
		fmt.Println(blue + msg + "\nRunning as non-root user" + reset)
		return
	}

	fmt.Println(redB + "This program must not be run as the root user" + reset)
	fmt.Println("")
	os.Exit(1)
}