package services

import (
	"context"
	"fmt"
	"time"

	"github.com/IBM/sarama"
	"github.com/shadow1ng/fscan/common"
)

type KafkaPlugin struct {
	name  string
	ports []int
}

func NewKafkaPlugin() *KafkaPlugin {
	return &KafkaPlugin{
		name:  "kafka",
		ports: []int{9092, 9093, 9094},
	}
}

func (p *KafkaPlugin) GetName() string {
	return p.name
}

func (p *KafkaPlugin) GetPorts() []int {
	return p.ports
}

func (p *KafkaPlugin) Scan(ctx context.Context, info *common.HostInfo) *ScanResult {
	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)

	if common.DisableBrute {
		return p.identifyService(ctx, info)
	}

	credentials := GenerateCredentials("kafka")
	if len(credentials) == 0 {
		return &ScanResult{
			Success: false,
			Service: "kafka",
			Error:   fmt.Errorf("没有可用的测试凭据"),
		}
	}

	for _, cred := range credentials {
		if client := p.testCredential(ctx, info, cred); client != nil {
			client.Close()
			common.LogSuccess(fmt.Sprintf("Kafka %s %s:%s", target, cred.Username, cred.Password))
			return &ScanResult{
				Success:  true,
				Service:  "kafka",
				Username: cred.Username,
				Password: cred.Password,
			}
		}
	}

	return &ScanResult{
		Success: false,
		Service: "kafka",
		Error:   fmt.Errorf("未发现弱密码"),
	}
}


func (p *KafkaPlugin) testCredential(ctx context.Context, info *common.HostInfo, cred Credential) sarama.Client {
	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
	timeout := time.Duration(common.Timeout) * time.Second

	config := sarama.NewConfig()
	config.Net.DialTimeout = timeout
	config.Net.ReadTimeout = timeout
	config.Net.WriteTimeout = timeout
	config.Version = sarama.V2_0_0_0

	if cred.Username != "" || cred.Password != "" {
		config.Net.SASL.Enable = true
		config.Net.SASL.Mechanism = sarama.SASLTypePlaintext
		config.Net.SASL.User = cred.Username
		config.Net.SASL.Password = cred.Password
		config.Net.SASL.Handshake = true
	}

	brokers := []string{target}
	client, err := sarama.NewClient(brokers, config)
	if err != nil {
		return nil
	}
	return client
}



func (p *KafkaPlugin) identifyService(ctx context.Context, info *common.HostInfo) *ScanResult {
	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)

	emptyCred := Credential{Username: "", Password: ""}
	client := p.testCredential(ctx, info, emptyCred)
	if client == nil {
		return &ScanResult{
			Success: false,
			Service: "kafka",
			Error:   fmt.Errorf("无法连接到Kafka服务"),
		}
	}
	defer client.Close()

	banner := "Kafka"
	common.LogSuccess(fmt.Sprintf("Kafka %s %s", target, banner))
	return &ScanResult{
		Success: true,
		Service: "kafka",
		Banner:  banner,
	}
}

func init() {
	// 使用高效注册方式：直接传递端口信息，避免实例创建
	RegisterPluginWithPorts("kafka", func() Plugin {
		return NewKafkaPlugin()
	}, []int{9092, 9093, 9094})
}