styx/plugins/shodan.go

package plugins

import (
	"context"
	"encoding/json"
	"sync"

	"github.com/dgraph-io/dgo/v2"
	"github.com/dgraph-io/dgo/v2/protos/api"
	"github.com/ns3777k/go-shodan/v4/shodan"
	"github.com/sirupsen/logrus"
	"github.com/spf13/viper"
	"gitlab.dcso.lolcat/LABS/styx/filters"
	"gitlab.dcso.lolcat/LABS/styx/models"
)

// ShodanPlugin defines the general ShodanPlugin structure.
type ShodanPlugin struct {
	Client      *shodan.Client
	ShodanChan  chan *shodan.HostData
	StopChan    chan bool
	StoppedChan chan bool
	Running     bool
}

// Initialize initialises the certstream configuration.
func (s *ShodanPlugin) Initialize() bool {
	if !viper.GetBool("shodan.activated") {
		return false
	}
	logrus.Info("shodan plugin is activated")
	s.ShodanChan = make(chan *shodan.HostData)
	return true
}

// Run runs the Shodan plugin.
func (s *ShodanPlugin) Run(wg *sync.WaitGroup, dgraphClient *dgo.Dgraph) {
	if !s.Running {
		s.StoppedChan = make(chan bool)
		wg.Add(1)
		go s.doRun(dgraphClient)
		s.Running = true
	}
}

// Stop stops the Shodan plugin.
func (s *ShodanPlugin) Stop(wg *sync.WaitGroup) {
	if s.Running {
		s.StoppedChan = make(chan bool)
		close(s.StopChan)
		<-s.StopChan
		wg.Done()
		s.Running = false
	}
}

func (s *ShodanPlugin) doRun(graphClient *dgo.Dgraph) {
	client := shodan.NewEnvClient(nil)
	err := client.GetBannersByPorts(context.Background(), viper.GetIntSlice("shodan.ports"), s.ShodanChan)
	if err != nil {
		logrus.Panic(err)
	}

	for {
		select {
		default:
			banner, ok := <-s.ShodanChan
			if !ok {
				logrus.Error("channel error: ", ok)
				break
			}

			shodanNode := models.BuildShodanNode(banner)
			// first filter poc
			if shodanNode.HostData.HTML != "" {
				if !filters.RunIPFilters(shodanNode.HostData.IP) {
					hostnames := shodanNode.HostData.Hostnames
					var hostNotInFilters, domainNotInFilters bool
					if len(hostnames) != 0 {
						for _, hostname := range hostnames {
							hostNotInFilters = filters.RunDomainFilters(hostname)
							if hostNotInFilters {
								logrus.Info("host", hostname, " not in filters")
								// keep track of new hostnames
								// saveSingleValues(conn, "shodan_stream", "hostname", shodanNode.ID, hostname)
							}
						}
					}
					domains := shodanNode.HostData.Domains
					if len(domains) != 0 {
						for _, domain := range domains {
							domainNotInFilters = filters.RunDomainFilters(domain)
							logrus.Info("domain", domain, "not in filters")
							// keep trakc of new domains
							// saveSingleValues(conn, "shodan_stream", "domain", shodanNode.ID, domain)
						}
					}
					if domainNotInFilters && hostNotInFilters {
						// models.SaveShodanNode("raw_shodan.json", shodanNode)
						mainNode := models.BuildNode("shodan", "shodan_stream", shodanNode.ID)
						models.SaveNode("nodes.json", mainNode)
						// edge := models.BuildEdge("shodan", structs.Map(shodanNode), structs.Map(mainNode))
						// models.SaveEdge(edge)
						e := models.Node{
							ID:         mainNode.ID,
							NodeType:   mainNode.NodeType,
							NData:      mainNode.NData,
							Created:    mainNode.Created,
							Modified:   mainNode.Modified,
							ShodanNode: *shodanNode,
						}

						ctx := context.Background()
						mu := &api.Mutation{
							CommitNow: true,
						}

						pb, err := json.Marshal(e)
						if err != nil {
							logrus.Fatal(err)
						}

						mu.SetJson = pb

						_, err = graphClient.NewTxn().Mutate(ctx, mu)
						if err != nil {
							logrus.Error(err)
						}
					}
				} else {
					logrus.Info(shodanNode.HostData.IP, "is akamain")
				}
			}
		}

	}
}