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 is closed")
				break
			}

			shodanNode := models.BuildShodanNode(banner)
			// first filter poc
			// if shodanNode.Data.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,
						Type:       mainNode.Type,
						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.Fatal(err)
					}
				}
			} else {
				logrus.Info(shodanNode.HostData.IP, "is akamain")
			}
		}
		// }

	}
}