styx/main.go

package main

import (
	"os"
	"sync"

	"github.com/sirupsen/logrus"
	"github.com/spf13/viper"
	"gitlab.dcso.lolcat/LABS/styx/broker"
	"gitlab.dcso.lolcat/LABS/styx/graph"
	"gitlab.dcso.lolcat/LABS/styx/matcher"
	"gitlab.dcso.lolcat/LABS/styx/plugins"
)

func init() {
	// Setting up logging.
	logrus.SetFormatter(&logrus.TextFormatter{
		FullTimestamp: true,
	})
	logrus.SetReportCaller(true)
	viper.SetConfigName("config")
	viper.SetConfigType("yaml")
	viper.AddConfigPath(".")
}

func main() {
	err := viper.ReadInConfig()
	if err != nil {
		panic(err)
	}

	os.Setenv("SHODAN_KEY", viper.GetString("shodan.key"))
	logrus.Info("Starting to get data from the Internet...")

	logrus.Info("Initializing Kafka...")
	_, err = broker.SetUpKafkaConnecter()
	if err != nil {
		logrus.WithField("err", err).Error("error initialising kafka")
	}
	logrus.Info("done")

	logrus.Info("Initializing Dgraph...")
	dgraphClient, err := graph.ConnectToDgraph()
	if err != nil {
		logrus.WithField("err", err).Error("error initialising the graph database")
	}
	logrus.Info("done")

	// stop channel
	stopChan := make(chan os.Signal)
	var wg sync.WaitGroup

	// go broker.ReadEventFromKafka()

	// certstream
	c := plugins.CertStreamPlugin{}
	if ok := c.Initialize(); !ok {
		logrus.Info("certstream plugin not activated")
	} else {
		c.Run(&wg, dgraphClient)
	}

	// pastebin
	p := plugins.PastebinPlugin{}
	if ok := p.Initialize(); !ok {
		logrus.Info("pastebin plugin not activated")
	} else {
		p.Run(&wg, dgraphClient)
	}

	// shodan
	s := plugins.ShodanPlugin{}
	if ok := s.Initialize(); !ok {
		logrus.Info("shodan plugin not activated")
	} else {
		s.Run(&wg, dgraphClient)
	}

	// main matcher
	m := matcher.Matcher{}
	if ok := m.Initialize(); !ok {
		logrus.Info("matcher is not activated")
	} else {
		m.Run(&wg, dgraphClient)
	}

	go func() {
		<-stopChan
		logrus.Info("Shutting down...")
		c.Stop(&wg)
		p.Stop(&wg)
		s.Stop(&wg)
		m.Stop(&wg)
	}()

	wg.Wait()
	logrus.Info("done")

}