package correlator

import (
	"encoding/json"
	"fmt"
	"github.com/google/uuid"
	"github.com/olivere/elastic/v7"
	log "github.com/sirupsen/logrus"
	"github.com/spf13/viper"
	"iwarma.ru/console/correlator/config"
	"iwarma.ru/console/correlator/es"
	"iwarma.ru/console/correlator/events"
	"iwarma.ru/console/correlator/rules"
	"iwarma.ru/console/correlator/util"
	"testing"
	"time"
)

// Check that one iteration of correlator works
func TestCorrelator(t *testing.T) {
	util.SetupTest(t)
	defer util.TearDownTest(t)

	client, err := es.NewElastic()
	if err != nil {
		t.Errorf("%v", err)
		return
	}

	err = events.ClearIndex(client, events.GetAggregatedIndexName(), events.GetNormalizedIndexName())
	if err != nil {
		t.Errorf("%v", err)
		return
	}

	// Events count
	N := 10

	err = events.FillAggregatedEvents(events.GetAggregatedIndexName(), N, client)
	if err != nil {
		t.Errorf("%v", err)
		return
	}

	// Prepare store
	store := rules.NewRuleStore(client)

	actions := make([]rules.Action, 1)
	action := rules.TestAction{
		PerformError:        false,
		ParseInterfaceError: false,
		MarshalError:        false,
		UnmarshalError:      false,
	}
	actions[0] = &action

	store.AddRules(rules.Rule{
		Id:          "1",
		Name:        "TestRule",
		Description: "Some long description",
		Multi:       false,
		Depth:       time.Second * 30,
		Predicate:   rules.NewPredicate("", "device_vendor:\"TestDevice\""),
		Actions:     actions,
	})

	rules.PerformCountClear()
	rules.EventsProcessedClear()

	// Prepare correlator
	stop := make(chan interface{})

	// Run correlator
	go func() {
		Correlator(store, stop)
	}()

	time.Sleep(time.Second * 2)

	stop <- ""

	// Let's check what we have
	goodCount := 1
	if rules.GetPerformCount() != goodCount {
		t.Errorf("Bad perform count. Expect %v, got %v", goodCount, rules.GetPerformCount())
	}

	if rules.GetEventsProcessed() != goodCount {
		t.Errorf("Bad events processed count. Expect %v, got %v", goodCount, rules.GetEventsProcessed())
	}

	stat := store.GetStat()

	if stat == nil {
		t.Errorf("No stat")
		return
	}

	if stat.IncidentCount != uint64(goodCount) {
		t.Errorf("Got bad stat, expect incident count eq %v, got %v", goodCount, stat.IncidentCount)
	}
}

// Check that several iterations of correlator work
func TestCorrelatorCycle(t *testing.T) {
	util.SetupTest(t)
	defer util.TearDownTest(t)

	client, err := es.NewElastic()
	if err != nil {
		t.Errorf("%v", err)
		return
	}

	err = events.ClearIndex(client, events.GetAggregatedIndexName(), events.GetNormalizedIndexName())
	if err != nil {
		t.Errorf("%v", err)
		return
	}

	// Prepare store
	store := rules.NewRuleStore(client)

	actions := make([]rules.Action, 1)
	action := rules.TestAction{
		PerformError:        false,
		ParseInterfaceError: false,
		MarshalError:        false,
		UnmarshalError:      false,
	}
	actions[0] = &action

	store.AddRules(rules.Rule{
		Id:          "1",
		Name:        "TestRule",
		Description: "Some long description",
		Multi:       false,
		Depth:       time.Second * 30,
		Predicate:   rules.NewPredicate("", "device_vendor:\"TestDevice\""),
		Actions:     actions,
	})

	rules.PerformCountClear()
	rules.EventsProcessedClear()

	// Prepare correlator
	stop := make(chan interface{})

	// Run correlator
	go func() {
		Correlator(store, stop)
	}()

	// Iteration count and event count in every iteration
	N := 10

	for i := 0; i < N; i++ {
		err = events.FillAggregatedEvents(events.GetAggregatedIndexName(), N, client)
		if err != nil {
			t.Errorf("%v", err)
			return
		}
	}

	// Here we need to wait until all events were send to elastic and correlator process them
	time.Sleep(time.Second * 4)
	stop <- ""
	time.Sleep(time.Second * 3)

	// Let's check what we have
	// No checks for GetPerformCount or EventsProcessed from action
	// Because we can't control how many times correlator will get the events

	stat := store.GetStat()

	if stat == nil {
		t.Errorf("No stat")
		return
	}

	if stat.EventsCount != uint64(N*N) {
		t.Errorf("Bad total events count, expect %v, got %v", N*N, stat.EventsCount)
	}

	log.Infof("Incident count: %v Events count: %v", rules.GetPerformCount(), rules.GetEventsProcessed())
}

// Check how correlator work in cycle with multi rule selected
func TestCorrelatorCycleMultiRule(t *testing.T) {
	util.SetupTest(t)
	defer util.TearDownTest(t)

	client, err := es.NewElastic()
	if err != nil {
		t.Errorf("%v", err)
		return
	}

	err = events.ClearIndex(client, events.GetAggregatedIndexName(), events.GetNormalizedIndexName())
	if err != nil {
		t.Errorf("%v", err)
		return
	}

	// Prepare store
	store := rules.NewRuleStore(client)

	actions := make([]rules.Action, 1)
	action := rules.TestAction{
		PerformError:        false,
		ParseInterfaceError: false,
		MarshalError:        false,
		UnmarshalError:      false,
	}
	actions[0] = &action

	store.AddRules(rules.Rule{
		Id:          "1",
		Name:        "TestRule",
		Description: "Some long description",
		Multi:       true,
		Depth:       time.Second * 30,
		Predicate:   rules.NewPredicate("", "device_vendor:\"TestDevice\""),
		Actions:     actions,
	})

	rules.PerformCountClear()
	rules.EventsProcessedClear()

	// Prepare correlator
	stop := make(chan interface{})

	// Run correlator
	go func() {
		Correlator(store, stop)
	}()

	// Iteration count and event count in every iteration
	N := 10

	for i := 0; i < N; i++ {
		err = events.FillAggregatedEvents(events.GetAggregatedIndexName(), N, client)
		if err != nil {
			t.Errorf("%v", err)
			return
		}
	}

	// Here we need to wait until all events were send to elastic and correlator process them
	time.Sleep(time.Second * 4)
	stop <- ""
	time.Sleep(time.Second * 3)

	// Let's check what we have
	// No checks for GetPerformCount or EventsProcessed from action
	// Because we can't control how many times correlator wil get events
	stat := store.GetStat()

	if stat == nil {
		t.Errorf("No stat")
		return
	}

	if stat.EventsCount != uint64(N*N) {
		t.Errorf("Bad total events count, expect %v, got %v", N*N, stat.EventsCount)
	}

	log.Infof("Incident count: %v Events count: %v", rules.GetPerformCount(), rules.GetEventsProcessed())
}

// Check how many incidents will we have
func TestCorrelatorIncidents(t *testing.T) {
	util.SetupTest(t)
	defer util.TearDownTest(t)

	client, err := es.NewElastic()
	if err != nil {
		t.Errorf("%v", err)
		return
	}

	err = events.ClearIndex(client, events.GetAggregatedIndexName(), events.GetNormalizedIndexName())
	if err != nil {
		t.Errorf("%v", err)
		return
	}

	// Prepare store
	store := rules.NewRuleStore(client)

	actions := make([]rules.Action, 1)
	action := rules.TestAction{
		PerformError:        false,
		ParseInterfaceError: false,
		MarshalError:        false,
		UnmarshalError:      false,
	}
	actions[0] = &action

	store.AddRules(rules.Rule{
		Id:          "1",
		Name:        "TestRule",
		Description: "Some long description",
		Multi:       false,
		Depth:       time.Second * 30,
		Predicate:   rules.NewPredicate("", "device_vendor:\"TestDevice\""),
		Actions:     actions,
	})

	rules.PerformCountClear()
	rules.EventsProcessedClear()

	// Prepare correlator
	stop := make(chan interface{})

	// Run correlator
	go func() {
		Correlator(store, stop)
	}()

	// Iteration count and event count in every iteration
	N := 10
	EventCount := 3

	for i := 0; i < EventCount; i++ {
		err = events.FillAggregatedEvents(events.GetAggregatedIndexName(), N, client)
		if err != nil {
			t.Errorf("%v", err)
			return
		}

		time.Sleep(time.Second * 3)
	}

	// Wait till correlator process all events
	time.Sleep(time.Second * 3)
	stop <- ""
	time.Sleep(time.Second * 3) // Wait for stop

	// Let's check what we have
	if rules.GetPerformCount() != EventCount {
		t.Errorf("Bad incident count. Expect %v got %v", EventCount, rules.GetPerformCount())
	}

	// Because of Multi = false in rule, we must have events count == incident count
	if rules.GetEventsProcessed() != EventCount {
		t.Errorf("Bad events count, expect %v got %v", EventCount, rules.GetEventsProcessed())
	}
}

func TestCorrelatorNoIndex(t *testing.T) {
	util.SetupTest(t)
	defer util.TearDownTest(t)

	client, err := es.NewElastic()
	if err != nil {
		t.Errorf("%v", err)
		return
	}

	err = events.ClearIndex(client, events.GetAggregatedIndexName(), events.GetNormalizedIndexName())
	if err != nil {
		t.Errorf("%v", err)
		return
	}

	// Prepare store
	store := rules.NewRuleStore(client)

	actions := make([]rules.Action, 1)
	action := rules.TestAction{
		PerformError:        false,
		ParseInterfaceError: false,
		MarshalError:        false,
		UnmarshalError:      false,
	}
	actions[0] = &action

	store.AddRules(rules.Rule{
		Id:          "1",
		Name:        "TestRule",
		Description: "Some long description",
		Multi:       false,
		Depth:       time.Second * 30,
		Predicate:   rules.NewPredicate("", "device_vendor:\"TestDevice\""),
		Actions:     actions,
	})

	rules.PerformCountClear()
	rules.EventsProcessedClear()

	// Prepare correlator
	stop := make(chan interface{})

	// Run correlator
	go func() {
		Correlator(store, stop)
	}()

	time.Sleep(time.Second * 5)
	stop <- ""
	time.Sleep(time.Second * 3)

	// Let's check what we have
	stat := store.GetStat()

	if stat.IncidentCount != 0 {
		t.Errorf("Bad incident count. Expect 0 got %v", stat.IncidentCount)
	}

	if stat.EventsCount != 0 {
		t.Errorf("Bad event count. Expect 0 got %v", stat.EventsCount)
	}
}

func TestCorrelatorActionError(t *testing.T) {
	util.SetupTest(t)
	defer util.TearDownTest(t)

	client, err := es.NewElastic()
	if err != nil {
		t.Errorf("%v", err)
		return
	}

	err = events.ClearIndex(client, events.GetAggregatedIndexName(), events.GetNormalizedIndexName())
	if err != nil {
		t.Errorf("%v", err)
		return
	}

	// Prepare store
	store := rules.NewRuleStore(client)

	actions := make([]rules.Action, 1)
	action := rules.TestAction{
		PerformError:        true,
		ParseInterfaceError: false,
		MarshalError:        false,
		UnmarshalError:      false,
	}
	actions[0] = &action

	store.AddRules(rules.Rule{
		Id:          "1",
		Name:        "TestRule",
		Description: "Some long description",
		Multi:       false,
		Depth:       time.Second * 30,
		Predicate:   rules.NewPredicate("", "device_vendor:\"TestDevice\""),
		Actions:     actions,
	})

	rules.PerformCountClear()
	rules.EventsProcessedClear()

	// Prepare correlator
	stop := make(chan interface{})

	// Run correlator
	go func() {
		Correlator(store, stop)
	}()

	// Iteration count and event count in every iteration
	N := 10
	EventCount := 3

	for i := 0; i < EventCount; i++ {
		err = events.FillAggregatedEvents(events.GetAggregatedIndexName(), N, client)
		if err != nil {
			t.Errorf("%v", err)
			return
		}

		time.Sleep(time.Second * 3)
	}

	// Wait till correlator process all events
	//time.Sleep(time.Second*3)
	stop <- ""
	time.Sleep(time.Second * 3) // Wait for stop

	stat := store.GetStat()
	if stat.IncidentCount != 0 {
		t.Errorf("Bad incident count #1. Expect 0 got %v", stat.IncidentCount)
	}

	goodEventCount := 1
	// Here we don't have events in store stat
	// because we count hem if no error occurs
	if stat.EventsCount != uint64(0) {
		t.Errorf("Bad event count #1. Expect %v got %v", 0, stat.EventsCount)
	}

	if rules.GetEventsProcessed() != goodEventCount {
		t.Errorf("Bad event count #2. Expect %v got %v", goodEventCount, rules.GetEventsProcessed())
	}

	if rules.GetPerformCount() != goodEventCount {
		t.Errorf("Bad incident count #2. Expect %v got %v", goodEventCount, rules.GetPerformCount())
	}

	log.Infof("%v", stat.String())
}

func TestRequest1(t *testing.T) {
	util.SetupTest(t)
	defer util.TearDownTest(t)

	viper.Set(config.ElasticLogQuery, true)

	client, err := es.NewElastic()
	if err != nil {
		t.Errorf("%v", err)
		return
	}

	err = events.ClearIndex(client, events.GetAggregatedIndexName(), events.GetNormalizedIndexName())
	if err != nil {
		t.Errorf("%v", err)
		return
	}

	// Let's prepare input
	const total = 150
	const good = 100

	var goodSid = [...]string{"act=2002910", "act=2001219", "act=2010937", "act=2009582", "act=2000538", "act=2010939", "act=2010935"}
	goodCounter := 0

	bulk := client.NewBulkRequest()

	for i := 0; i < total; i++ {

		srcMsg := "CEF:0|InfoWatch ARMA|ARMAIF|3.0|idspower|IDS power|5|rt=1604793739000 log_from=suricata deviceFacility=28775 gid=1 cs1=429496728 cs1Label=Signature cs2=12 cs2Label=line_number rev=1 classification=null priority=3 proto=TCP ip_src=10.20.30.50 port_src=80 ip_dst=10.20.30.1 port_dst=34568 %v"
		if goodCounter < good {
			srcMsg = fmt.Sprintf(srcMsg, goodSid[goodCounter%len(goodSid)])
			goodCounter += 1
		} else {
			srcMsg = fmt.Sprintf(srcMsg, "act=6000")
		}

		event := events.Event{
			events.AggregatedId: uuid.NewString(),
			events.FirstEvent:   time.Now().UTC(),
			events.LastEvent:    time.Now().UTC(),
			events.EventCount:   total,
			events.Created:      time.Now().UTC(),
			events.Tags:         make([]string, 0),
			events.Hash:         uuid.NewString(),
			events.CeleryDone:   false,
			"timestamp":         time.Now().UTC(),
			"type":              "ass",
			"index":             events.GetAggregatedIndexName(),
			"event_timestamp":   time.Now().UTC(),
			"event_id":          uuid.NewString(),
			"event_severity":    uint8(i % 10),
			"event_src_msg":     srcMsg,
			"event_protocol":    "TCP",
			"device_vendor":     "TestDevice",
			"device_product":    "TestProduct",
			"device_version":    "1.0",
			"device_action":     "Test",
			"sign_ad":           fmt.Sprintf("%v", i),
			"sign_category":     "Test",
			"sign_subcategory":  "Test",
			"application":       "Test",
			"source_ip":         "127.0.0.1",
			"source_mac":        "00:50:56:c0:00:08",
			"source_host":       "localhost",
			"source_port":       uint32(i),
			"source_user":       "root",
			"destination_ip":    "127.0.0.1",
			"destination_host":  "localhost",
			"destination_port":  uint32(i),
			"destination_user":  "user",
		}

		bulk = bulk.Add(elastic.NewBulkIndexRequest().Index(events.GetAggregatedIndexName()).Id(event.GetString(events.EventID)).Doc(event))
	}

	bulkResponse, err := client.ExecuteBulk(bulk)
	if err != nil {
		t.Errorf("Can't index documents: %v", err)
		return
	}

	if bulkResponse.Errors {
		t.Errorf("Got errors from bulk requset: %v", bulkResponse.Failed())
		return
	}

	if len(bulkResponse.Indexed()) != total {
		t.Errorf("Bad bulk index count. Got %v, expect %v", len(bulkResponse.Indexed()), total)
		return
	}

	// wait until elastic is ready

	time.Sleep(time.Second)

	// Now we need construct rule
	ruleStr := `{
   "name":"Test",
   "depth":"600s",
   "id":"1",
   "predicat":{
      "type":"query_string",
      "field":"NULL",
      "value":[
         "event_src_msg: \"act=2002910\" OR event_src_msg: \"act=2001219\" OR event_src_msg: \"act=2010937\" OR event_src_msg: \"act=2009582\" OR event_src_msg: \"act=2000538\" OR event_src_msg: \"act=2010939\" OR event_src_msg: \"act=2010935\""
      ]
   },
   "actions":[
      {
         "host":"localhost",
         "port":"514",
         "protocol":"udp",
         "name":"test",
         "template":"{{.event_src_msg}}",
         "type":"syslog"
      }
   ]
}`
	var rule rules.Rule
	err = json.Unmarshal([]byte(ruleStr), &rule)
	if err != nil {
		t.Errorf("Can't read rule: %v", err)
		return
	}

	// Now, let's run it
	eventList, err := rule.Do(client)
	if err != nil {
		t.Errorf("Can't do rule: %v", err)
		return
	}

	if len(*eventList) != good {
		t.Errorf("Bad event's count. Expect %v, got %v", good, len(*eventList))
	}
}
