log_consumer.go

package service

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"time"

	"github.com/rs/zerolog/log"
	"github.com/segmentio/kafka-go"
	"github.com/segmentio/kafka-go/sasl/scram"
	"gitlab.com/tensorsecurity-rd/waf-console/internal/config"
	"gitlab.com/tensorsecurity-rd/waf-console/internal/model"
	"gitlab.com/tensorsecurity-rd/waf-console/internal/utils/id"

	es "github.com/olivere/elastic/v7"
	esStore "gitlab.com/tensorsecurity-rd/waf-console/internal/store"
	"gorm.io/gorm"
)

const (
	KafkaTopicIvanWafDetections       = "ivan_waf_detections"
	KafkaGroupIvanWafDetectionsPalace = "ivan_waf_detections_palace"

	EsIndexWafDetections      = "waf-detections*"
	EsIndexWafDetectionsAlias = "waf-detections"
	ESIndexEvents             = "events"
)

var scramAlgo = map[string]scram.Algorithm{
	"SHA512": scram.SHA512,
	"SHA256": scram.SHA256,
}

var ErrMissingBrokers = errors.New("no KAFKA_BROKERS env variable set")
var ErrMissingAuthInfo = errors.New("no username or passwrod set when auth is enabled")
var ErrUnsupportedAuthMethod = errors.New("unsupported auth method")

const (
	EnvKafkaAuthMethod   = "KAFKA_AUTH_METHOD"
	EnvKafkaAuthUsername = "KAFKA_AUTH_USERNAME"
	EnvKafkaAuthPassword = "KAFKA_AUTH_PASSWORD"
	EnvKafkaScramAlgo    = "KAFKA_SCRAM_ALGO"
	EnvKafkaBrokers      = "KAFKA_BROKERS"
	KafkaAuthPlain       = "plain"
	KafkaAuthScram       = "scram"
)

// func getSASLMechanismByEnv() (sasl.Mechanism, bool, error) {
// 	authMethod := os.Getenv(EnvKafkaAuthMethod)
// 	username := os.Getenv(EnvKafkaAuthUsername)
// 	password := os.Getenv(EnvKafkaAuthPassword)
// 	var mechanism sasl.Mechanism
// 	switch authMethod {
// 	case KafkaAuthPlain:
// 		if username == "" || password == "" {
// 			return nil, true, ErrMissingAuthInfo
// 		}
// 		mechanism = &plain.Mechanism{
// 			Username: username,
// 			Password: password,
// 		}
// 	case KafkaAuthScram:
// 		if username == "" || password == "" {
// 			return nil, true, ErrMissingAuthInfo
// 		}
// 		algoKey := os.Getenv(EnvKafkaScramAlgo)
// 		algo := scram.SHA512
// 		algoTmp, exist := scramAlgo[algoKey]
// 		if exist {
// 			algo = algoTmp
// 		}
// 		var err error
// 		mechanism, err = scram.Mechanism(algo, username, password)
// 		if err != nil {
// 			return nil, true, err
// 		}
// 	}
// 	return mechanism, mechanism != nil, nil
// }

type LogConsumerService struct {
	consumer *kafka.Reader
	db       *gorm.DB
	esStore  *esStore.ESStore
}

func NewLogConsumerService(db *gorm.DB, esStore *esStore.ESStore, config *config.KafkaConfig) *LogConsumerService {
	// mechanism, _, err := getSASLMechanismByEnv()
	// if err != nil {
	// 	log.Fatal().Err(err).Msg("failed to get SASL mechanism")
	// }
	// brokers := strings.Split(os.Getenv(EnvKafkaBrokers), ",")
	// if len(brokers) == 0 {
	// 	log.Fatal().Msg("no KAFKA_BROKERS env variable set")
	// }

	// consumer := kafka.NewReader(kafka.ReaderConfig{
	// 	Brokers: brokers,
	// 	GroupID: KafkaGroupIvanWafDetectionsPalace,
	// 	Topic:   KafkaTopicIvanWafDetections,
	// 	Dialer: &kafka.Dialer{
	// 		SASLMechanism: mechanism,
	// 		Timeout:       10 * time.Second,
	// 		DualStack:     true,
	// 	},
	// })
	consumer := NewKafkaReader(config)

	return &LogConsumerService{
		consumer: consumer,
		db:       db,
		esStore:  esStore,
	}
}

func (s *LogConsumerService) Consume() {
	log.Info().Msg("start consume kafka message")
	for {
		m, err := s.consumer.ReadMessage(context.Background())
		if err != nil {
			log.Error().Err(err).Msg("failed to read message")
			continue
		}
		log.Info().Msgf("message: %s", string(m.Value))
		go s.processMessage(m)
	}
}

func (s *LogConsumerService) processMessage(m kafka.Message) {
	log.Info().Msgf("processing message: %s", string(m.Value))
	s.Handle(context.Background(), m.Value)
}

func (s *LogConsumerService) genWafDetection(wafDetectionMessage model.WafDetectionMessage, attackedLog model.WafDetectionAttackedLog) (model.WafDetection, error) {
	if attackedLog.AttackIP == "" {
		return model.WafDetection{}, errors.New("attack_ip is empty")
	}

	wafDetection := model.WafDetection{
		WafDetectionMessageBasic: wafDetectionMessage.WafDetectionMessageBasic,
		WafDetectionAttackedLog:  attackedLog,
		CreatedAt:                wafDetectionMessage.CreatedAt,
	}
	wafDetection.WafDetectionAttackedLog.ID = id.Str()
	return wafDetection, nil
}

func (s *LogConsumerService) genWafDetectionEvent(wafDetectionMessage model.WafDetectionMessage, attackedLog model.WafDetectionAttackedLog) (model.Event, error) {
	event := model.Event{
		ID:          id.Str(),
		Type:        "waf_detection",
		Description: "waf detection",
		RuleKeys:    []model.RuleKey{},
		Scopes: map[string][]model.Scope{
			"cluster": {
				{
					Kind: "cluster",
					ID:   wafDetectionMessage.WafDetectionMessageBasic.ClusterKey,
					Name: wafDetectionMessage.WafDetectionMessageBasic.ClusterKey,
				},
			},
			"namespace": {
				{
					Kind: "namespace",
					ID:   wafDetectionMessage.WafDetectionMessageBasic.Namespace,
					Name: wafDetectionMessage.WafDetectionMessageBasic.Namespace,
				},
			},
			"resource": {
				{
					Kind: "resource",
					ID:   "",
					Name: fmt.Sprintf("%s(%s)", wafDetectionMessage.WafDetectionMessageBasic.ResName, wafDetectionMessage.WafDetectionMessageBasic.ResKind),
				},
			},
		},
		Resources: []map[string]model.Scope{
			{
				"cluster": {
					Kind: "cluster",
					ID:   wafDetectionMessage.WafDetectionMessageBasic.ClusterKey,
					Name: wafDetectionMessage.WafDetectionMessageBasic.ClusterKey,
				},
				"namespace": {
					Kind: "namespace",
					ID:   wafDetectionMessage.WafDetectionMessageBasic.Namespace,
					Name: wafDetectionMessage.WafDetectionMessageBasic.Namespace,
				},
				"resource": {
					Kind: "resource",
					ID:   "",
					Name: fmt.Sprintf("%s(%s)", wafDetectionMessage.WafDetectionMessageBasic.ResName, wafDetectionMessage.WafDetectionMessageBasic.ResKind),
				},
			},
		},
		Relation: model.Relation{
			Type: "timeline",
		},
		CreatedAt: wafDetectionMessage.CreatedAt,
		UpdatedAt: wafDetectionMessage.CreatedAt,
		Timestamp: time.Now(),
		Context: map[string]interface{}{
			"attack_ip":   attackedLog.AttackIP,
			"attack_time": attackedLog.AttackTime,
			"attack_url":  attackedLog.AttackedURL,
			"attack_app":  attackedLog.AttackedApp,
			"attack_load": attackedLog.AttackLoad,
			"rule_name":   attackedLog.RuleName,
			"action":      attackedLog.Action,
			"waf_body": map[string]interface{}{
				"type":     "code",
				"request":  attackedLog.ReqPkg,
				"response": attackedLog.RspPkg,
			},
		},
	}
	return event, nil
}

func (s *LogConsumerService) Handle(ctx context.Context, message []byte) error {
	WafDetectionMessage := model.WafDetectionMessage{}
	err := json.Unmarshal(message, &WafDetectionMessage)
	if err != nil {
		log.Err(err).Str("message.Value", string(message)).Msg("unmarshal kafka message fails")
		return err
	}

	bulkableRequests := make([]es.BulkableRequest, 0)
	// WafDetections := make([]model.WafDetection, len(WafDetectionMessage.AttackedLog))

	unPassCount := 0
	for i := range WafDetectionMessage.AttackedLog {
		if WafDetectionMessage.AttackedLog[i].AttackIP == "" {
			log.Err(err).Str("message.Value", string(message)).Msg("WafDetectionMessage.AttackedLog[i].AttackIP empty")
			continue
		}
		if WafDetectionMessage.AttackedLog[i].Action != "pass" {
			unPassCount++
		}
		// WafDetections[i].WafDetectionMessageBasic = WafDetectionMessage.WafDetectionMessageBasic
		// WafDetections[i].WafDetectionAttackedLog = WafDetectionMessage.AttackedLog[i]
		// WafDetections[i].WafDetectionAttackedLog.ID = id.Str()
		// WafDetections[i].CreatedAt = WafDetectionMessage.CreatedAt
		wafDetection, err := s.genWafDetection(WafDetectionMessage, WafDetectionMessage.AttackedLog[i])
		if err != nil {
			log.Err(err).Str("message.Value", string(message)).Msg("gen waf detection fails")
			continue
		}

		bulkIndexSignal := es.NewBulkIndexRequest().Index(EsIndexWafDetectionsAlias)
		bulkableRequests = append(bulkableRequests, bulkIndexSignal.Id(wafDetection.WafDetectionAttackedLog.ID).Doc(wafDetection))

		event, err := s.genWafDetectionEvent(WafDetectionMessage, WafDetectionMessage.AttackedLog[i])
		if err != nil {
			log.Err(err).Str("message.Value", string(message)).Msg("gen waf detection event fails")
			continue
		}
		bulkIndexEvent := es.NewBulkIndexRequest().Index(ESIndexEvents)
		bulkableRequests = append(bulkableRequests, bulkIndexEvent.Id(event.ID).Doc(event))
	}

	s.esStore.Save(ctx, bulkableRequests)

	// err = s.db.WithContext(ctx).Model(&model.WafService{}).Where("id = ?", WafDetectionMessage.ServiceID).Update("attack_num", gorm.Expr("attack_num + ?", 1)).Error
	// if err != nil {
	// 	log.Err(err).Int64("WafDetectionMessage.ServiceID", WafDetectionMessage.ServiceID).Msg("update waf_service attack_number fails")
	// 	return err
	// }

	return nil
}