package main

import (
	"context"
	"fmt"
	"os"
	"sync"
	"time"

	"github.com/CaliDog/certstream-go"
	"github.com/ns3777k/go-shodan/v4/shodan"
	"github.com/segmentio/kafka-go"
	"github.com/sirupsen/logrus"
	"gitlab.dcso.lolcat/LABS/styx/broker"
	"gitlab.dcso.lolcat/LABS/styx/models"
)

func init() {
	// Setting up logging.
	logrus.SetFormatter(&logrus.TextFormatter{
		FullTimestamp: true,
	})
	logrus.SetReportCaller(true)
}

func main() {
	fmt.Println("Starting to get data from the Internet...")

	// The false flag specifies that we want heartbeat messages.
	stream, errStream := certstream.CertStreamEventStream(false)
	Conn, err := broker.SetUpKafkaConnecter()
	if err != nil {
		panic(err)
	}

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

	go broker.ReadEventFromKafka()

	// certstream
	go func() {
		for {
			select {
			case jq := <-stream:
				if data, err := models.ExtractCertFromStream(jq); err == nil {

					rawNode := models.WrapCertStreamData(*data)
					models.SaveCertStreamRaw("raw_certstream.json", rawNode)

					certNode := models.BuildCertNode(rawNode)
					models.SaveCertNode("cert_nodes.json", certNode)
					mainNode := models.BuildNode("node", "certstream", certNode.ID)
					models.SaveNode("nodes.json", mainNode)
					edge := models.BuildEdge("certstream", rawNode.ID, mainNode.ID)
					models.SaveEdge(edge)
					edge = models.BuildEdge("certstream", mainNode.ID, certNode.ID)
					models.SaveEdge(edge)
					allDomains := data.Data.LeafCert.AllDomains
					saveSingleValues(Conn, "certstream", "domain", certNode.ID, allDomains)
				}

			case err := <-errStream:
				logrus.Error(err)
			case <-stopChan:
				wg.Done()
				return
			}
		}

	}()

	// pastebin
	go func() {
		for {
			select {
			default:
				pastes, err := models.QueryPastes()
				if err != nil {
					logrus.Panic(err)
				}
				for _, p := range pastes {
					paste, err := models.FetchPaste(p)
					if err != nil {
						logrus.Error("cannot fetch paste", err)
					}
					fp := models.FullPaste{
						Meta: p,
						Full: paste,
					}
					res := models.BuildPasteNode(&fp)
					models.SavePaste("paste_formatted.json", res)
					time.Sleep(1 * time.Second)

				}
				time.Sleep(3 * time.Second)
			case <-stopChan:
				wg.Done()
				return
			}
		}
	}()

	// shodan
	client := shodan.NewEnvClient(nil)
	ch := make(chan *shodan.HostData)
	err = client.GetBannersByPorts(context.Background(), []int{80, 443, 8443, 53}, ch)
	if err != nil {
		logrus.Panic(err)
	}

	go func() {
		for {
			select {
			default:
				banner, ok := <-ch
				if !ok {
					logrus.Error("channel is closed")
					break
				}

				shodanNode := models.BuildShodanNode(banner)
				hostnames := shodanNode.Data.Hostnames
				if len(hostnames) != 0 {
					saveSingleValues(Conn, "shodan_stream", "hostname", shodanNode.ID, hostnames)
				}
				domains := shodanNode.Data.Domains
				if len(domains) != 0 {
					saveSingleValues(Conn, "shodan_stream", "domain", shodanNode.ID, domains)
				}
				models.SaveShodanNode("shodan_raw.json", shodanNode)
				node := models.BuildNode("shodan", "shodan_stream", shodanNode.ID)
				models.SaveNode("nodes.json", node)
				edge := models.BuildEdge("shodan", shodanNode.ID, node.ID)
				models.SaveEdge(edge)
			case <-stopChan:
				wg.Done()
				return
			}
		}
	}()
	wg.Wait()

}

// helpers
func saveSingleValues(brokerConn *kafka.Conn, source string, datatype string, originNodeID string, values []string) {
	for _, value := range values {
		domainNode := models.BuildNode(source, datatype, value)
		models.SaveNode("nodes.json", domainNode)
		if domainNode.Type == "domain" || domainNode.Type == "hostname" {
			broker.SendEventToKafka(brokerConn, *domainNode)
		}
		edge := models.BuildEdge(source, originNodeID, domainNode.ID)
		models.SaveEdge(edge)
	}
}