alert_consumer.go

package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	"github.com/confluentinc/confluent-kafka-go/v2/kafka"
)

// AlertData struct để parse JSON từ Kafka
type AlertData struct {
	AlertID       string `json:"alert_id"`
	Timestamp     string `json:"timestamp"`
	DetectionTime string `json:"detection_time"`
	Severity      string `json:"severity"`
	Type          string `json:"type"`
	Source        struct {
		IP   string `json:"ip"`
		Port int    `json:"port"`
	} `json:"source"`
	Destination struct {
		IP   string `json:"ip"`
		Port int    `json:"port"`
	} `json:"destination"`
	Protocol   int `json:"protocol"`
	Prediction struct {
		Label      string  `json:"label"`
		Confidence float64 `json:"confidence"`
		Threshold  float64 `json:"threshold"`
	} `json:"prediction"`
	FlowStats struct {
		Duration         float64 `json:"duration"`
		TotalFwdPackets  int     `json:"total_fwd_packets"`
		TotalBwdPackets  int     `json:"total_bwd_packets"`
		TotalFwdBytes    int     `json:"total_fwd_bytes"`
		TotalBwdBytes    int     `json:"total_bwd_bytes"`
	} `json:"flow_stats"`
}

// getProtocolName converts protocol number to name
func getProtocolName(protocol int) string {
	protocolMap := map[int]string{
		1:  "ICMP",
		6:  "TCP",
		17: "UDP",
	}
	if name, ok := protocolMap[protocol]; ok {
		return name
	}
	return fmt.Sprintf("Protocol-%d", protocol)
}

// escapeHTML escapes special characters for Telegram HTML parse mode
func escapeHTML(s string) string {
	s = strings.ReplaceAll(s, "&", "&")
	s = strings.ReplaceAll(s, "<", "&lt;")
	s = strings.ReplaceAll(s, ">", "&gt;")
	return s
}

func sendTelegramAlert(message string) error {
	credentials := []string{
		"8322060808:AAFeGSVmuhGceQyJVkSKc746odROERyL0A0",
		"5833402318",
	}
	botToken := credentials[0]
	chatID := credentials[1]

	url := fmt.Sprintf("https://api.telegram.org/bot%s/sendMessage", botToken)
	
	// Escape message for JSON
	escapedMsg := strings.ReplaceAll(message, `"`, `\"`)
	escapedMsg = strings.ReplaceAll(escapedMsg, "\n", "\\n")
	
	bodyJSON := fmt.Sprintf(
		`{"chat_id":"%s","text":"%s","parse_mode":"HTML"}`,
		chatID, escapedMsg,
	)

	resp, err := http.Post(url, "application/json", bytes.NewBuffer([]byte(bodyJSON)))
	if err != nil {
		return fmt.Errorf("failed to send POST request: %w", err)
	}
	defer resp.Body.Close()

	// Check if response status is a success (2xx)
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		// Read the error body
		bodyBytes, readErr := io.ReadAll(resp.Body)
		if readErr != nil {
			return fmt.Errorf("status %s; error reading body: %w", resp.Status, readErr)
		}
		return fmt.Errorf("telegram API error: status %s; body: %s", resp.Status, string(bodyBytes))
	}

	// All good; success
	return nil
}

func main() {

	c, err := kafka.NewConsumer(&kafka.ConfigMap{
		"bootstrap.servers": "localhost:9092",
		"group.id":          "myGroup",
		"auto.offset.reset": "earliest",
	})

	if err != nil {
		panic(err)
	}

	err = c.SubscribeTopics([]string{"alert"}, nil)

	if err != nil {
		panic(err)
	}

	// A signal handler or similar could be used to set this to false to break the loop.
	run := true

	for run {
		msg, err := c.ReadMessage(time.Second)
		if err == nil {
			fmt.Printf("Message on %s: %s\n", msg.TopicPartition, string(msg.Value))
			
			// Parse JSON alert
			var alert AlertData
			if err := json.Unmarshal(msg.Value, &alert); err != nil {
				fmt.Printf("❌ Failed to parse alert JSON: %v\n", err)
				continue
			}
			
			// Format message với thông tin chi tiết
			telegramMsg := fmt.Sprintf(
				"🚨 <b>BOTNET ALERT</b> 🚨\n\n"+
				"<b>Alert ID:</b> %s\n"+
				"<b>Time:</b> %s\n"+
				"<b>Severity:</b> %s\n\n"+
				"<b>🔴 Detection:</b>\n"+
				"  • Label: %s\n"+
				"  • Confidence: %.2f%%\n"+
				"  • Threshold: %.4f\n\n"+
				"<b>📡 Network Info:</b>\n"+
				"  • Source: %s:%d\n"+
				"  • Destination: %s:%d\n"+
				"  • Protocol: %s\n\n"+
				"<b>📊 Flow Statistics:</b>\n"+
				"  • Duration: %.2f sec\n"+
				"  • Forward: %d packets (%d bytes)\n"+
				"  • Backward: %d packets (%d bytes)",
				escapeHTML(alert.AlertID),
				escapeHTML(alert.Timestamp),
				escapeHTML(alert.Severity),
				escapeHTML(alert.Prediction.Label),
				alert.Prediction.Confidence*100,
				alert.Prediction.Threshold,
				escapeHTML(alert.Source.IP),
				alert.Source.Port,
				escapeHTML(alert.Destination.IP),
				alert.Destination.Port,
				getProtocolName(alert.Protocol),
				alert.FlowStats.Duration,
				alert.FlowStats.TotalFwdPackets,
				alert.FlowStats.TotalFwdBytes,
				alert.FlowStats.TotalBwdPackets,
				alert.FlowStats.TotalBwdBytes,
			)
			
			// Gửi alert đến Telegram
			err := sendTelegramAlert(telegramMsg)
			if err != nil {
				fmt.Printf("❌ Failed to send Telegram alert: %v\n", err)
			} else {
				fmt.Println("✅ Alert sent to Telegram successfully")
			}

		} else if !err.(kafka.Error).IsTimeout() {
			// The client will automatically try to recover from all errors.
			// Timeout is not considered an error because it is raised by
			// ReadMessage in absence of messages.
			fmt.Printf("Consumer error: %v (%v)\n", err, msg)
		}
	}

	c.Close()
}