package plugins

import (
	"context"
	"fmt"
	"net"
	"strings"
	"time"

	"github.com/shadow1ng/fscan/common"
	"github.com/shadow1ng/fscan/common/i18n"
)

// ActiveMQPlugin ActiveMQ消息队列扫描插件 - 基于STOMP协议
type ActiveMQPlugin struct {
	name  string
	ports []int
}

// NewActiveMQPlugin 创建ActiveMQ插件
func NewActiveMQPlugin() *ActiveMQPlugin {
	return &ActiveMQPlugin{
		name:  "activemq",
		ports: []int{61616, 61617, 61618, 8161}, // STOMP端口 + Web管理界面
	}
}

// GetName 实现Plugin接口
func (p *ActiveMQPlugin) GetName() string {
	return p.name
}

// GetPorts 实现Plugin接口
func (p *ActiveMQPlugin) GetPorts() []int {
	return p.ports
}

// Scan 执行ActiveMQ扫描 - STOMP协议弱密码检测
func (p *ActiveMQPlugin) Scan(ctx context.Context, info *common.HostInfo) *ScanResult {
	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)

	// 如果禁用暴力破解，只做服务识别
	if common.DisableBrute {
		return p.identifyService(info)
	}

	// 生成测试凭据
	credentials := GenerateCredentials("activemq")
	if len(credentials) == 0 {
		// ActiveMQ默认凭据
		credentials = []Credential{
			{Username: "admin", Password: "admin"},
			{Username: "admin", Password: ""},
			{Username: "admin", Password: "password"},
			{Username: "user", Password: "user"},
			{Username: "guest", Password: "guest"},
		}
	}

	// 逐个测试凭据
	for _, cred := range credentials {
		// 检查Context是否被取消
		select {
		case <-ctx.Done():
			return &ScanResult{
				Success: false,
				Service: "activemq",
				Error:   ctx.Err(),
			}
		default:
		}

		// 测试凭据
		if p.testCredential(ctx, info, cred) {
			// ActiveMQ认证成功
			common.LogSuccess(i18n.GetText("activemq_scan_success", target, cred.Username, cred.Password))

			return &ScanResult{
				Success:  true,
				Service:  "activemq",
				Username: cred.Username,
				Password: cred.Password,
			}
		}
	}

	// 所有凭据都失败
	return &ScanResult{
		Success: false,
		Service: "activemq",
		Error:   fmt.Errorf("未发现弱密码"),
	}
}

// testCredential 测试单个凭据 - STOMP协议认证
func (p *ActiveMQPlugin) testCredential(ctx context.Context, info *common.HostInfo, cred Credential) bool {
	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
	timeout := time.Duration(common.Timeout) * time.Second

	// 建立TCP连接
	conn, err := common.WrapperTcpWithTimeout("tcp", target, timeout)
	if err != nil {
		return false
	}
	defer conn.Close()

	// 使用STOMP协议进行认证
	return p.authenticateSTOMP(conn, cred.Username, cred.Password)
}

// authenticateSTOMP 使用STOMP协议进行身份验证
// STOMP (Simple Text Oriented Messaging Protocol) 是ActiveMQ支持的文本协议
func (p *ActiveMQPlugin) authenticateSTOMP(conn net.Conn, username, password string) bool {
	timeout := time.Duration(common.Timeout) * time.Second

	// 构造STOMP CONNECT帧
	// STOMP是基于帧的协议，每个帧以NULL字符结尾
	stompConnect := fmt.Sprintf("CONNECT\naccept-version:1.0,1.1,1.2\nhost:/\nlogin:%s\npasscode:%s\n\n\x00",
		username, password)

	// 设置写超时并发送认证请求
	conn.SetWriteDeadline(time.Now().Add(timeout))
	if _, err := conn.Write([]byte(stompConnect)); err != nil {
		return false
	}

	// 设置读超时并读取响应
	conn.SetReadDeadline(time.Now().Add(timeout))
	response := make([]byte, 1024)
	n, err := conn.Read(response)
	if err != nil || n == 0 {
		return false
	}

	responseStr := string(response[:n])

	// 检查STOMP响应
	// 成功响应应该包含"CONNECTED"帧
	// 失败响应包含"ERROR"帧
	if strings.Contains(responseStr, "CONNECTED") {
		return true
	} else if strings.Contains(responseStr, "ERROR") {
		// 错误响应，认证失败
		return false
	}

	// 未知响应格式
	return false
}

// identifyService 服务识别 - 检测STOMP协议
func (p *ActiveMQPlugin) identifyService(info *common.HostInfo) *ScanResult {
	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
	timeout := time.Duration(common.Timeout) * time.Second

	// 尝试连接
	conn, err := common.WrapperTcpWithTimeout("tcp", target, timeout)
	if err != nil {
		return &ScanResult{
			Success: false,
			Service: "activemq",
			Error:   err,
		}
	}
	defer conn.Close()

	// 发送简单的STOMP CONNECT帧（无认证信息）
	stompConnect := "CONNECT\naccept-version:1.0,1.1,1.2\nhost:/\n\n\x00"
	
	conn.SetWriteDeadline(time.Now().Add(timeout))
	if _, err := conn.Write([]byte(stompConnect)); err != nil {
		return &ScanResult{
			Success: false,
			Service: "activemq",
			Error:   fmt.Errorf("无法发送STOMP请求: %v", err),
		}
	}

	// 读取响应
	conn.SetReadDeadline(time.Now().Add(timeout))
	response := make([]byte, 512)
	n, err := conn.Read(response)
	if err != nil || n == 0 {
		return &ScanResult{
			Success: false,
			Service: "activemq",
			Error:   fmt.Errorf("无法读取响应"),
		}
	}

	responseStr := string(response[:n])

	// 检查是否为STOMP协议响应
	if strings.Contains(responseStr, "CONNECTED") || strings.Contains(responseStr, "ERROR") {
		banner := p.extractSTOMPVersion(responseStr)
		common.LogSuccess(i18n.GetText("activemq_service_identified", target, banner))
		
		return &ScanResult{
			Success: true,
			Service: "activemq",
			Banner:  banner,
		}
	}

	return &ScanResult{
		Success: false,
		Service: "activemq",
		Error:   fmt.Errorf("无法识别为ActiveMQ STOMP服务"),
	}
}

// extractSTOMPVersion 从STOMP响应中提取版本信息
func (p *ActiveMQPlugin) extractSTOMPVersion(response string) string {
	lines := strings.Split(response, "\n")
	
	for _, line := range lines {
		// 查找version头
		if strings.HasPrefix(line, "version:") {
			version := strings.TrimPrefix(line, "version:")
			return fmt.Sprintf("ActiveMQ STOMP %s", version)
		}
		// 查找server头
		if strings.HasPrefix(line, "server:") {
			server := strings.TrimPrefix(line, "server:")
			return fmt.Sprintf("ActiveMQ %s", server)
		}
	}

	// 如果没有找到版本信息，返回通用描述
	if strings.Contains(response, "CONNECTED") {
		return "ActiveMQ STOMP (版本未知)"
	} else if strings.Contains(response, "ERROR") {
		return "ActiveMQ STOMP (需要认证)"
	}

	return "ActiveMQ STOMP"
}

// init 自动注册插件
func init() {
	RegisterPlugin("activemq", func() Plugin {
		return NewActiveMQPlugin()
	})
}