log_consumer.go

package service

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

	"github.com/rs/zerolog/log"
	"github.com/segmentio/kafka-go"
	"github.com/segmentio/kafka-go/sasl/scram"
	"gitlab.com/tensorsecurity-rd/waf-console/cmd/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"
)

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) 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

		bulkIndexSignal := es.NewBulkIndexRequest().Index(EsIndexWafDetectionsAlias)
		bulkableRequests = append(bulkableRequests, bulkIndexSignal.Id(WafDetections[i].WafDetectionAttackedLog.ID).Doc(WafDetections[i]))
	}

	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
}