package services

import (
	"context"
	"fmt"
	"io"
	"net"
	"net/http"
	"strings"
	"time"

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

type RabbitMQPlugin struct {
	plugins.BasePlugin
}

func NewRabbitMQPlugin() *RabbitMQPlugin {
	return &RabbitMQPlugin{
		BasePlugin: plugins.NewBasePlugin("rabbitmq"),
	}
}



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

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

	// 对于AMQP端口，首先识别服务
	if info.Ports == "5672" || info.Ports == "5671" {
		if result := p.testAMQPProtocol(ctx, info); result.Success {
			// AMQP协议识别成功，尝试HTTP管理接口的密码爆破
			managementResult := p.testManagementInterface(ctx, info)
			if managementResult.Success {
				return managementResult
			}
			// 返回服务识别结果
			return result
		}
	}

	// HTTP端口的密码爆破
	credentials := plugins.GenerateCredentials("rabbitmq")
	if len(credentials) == 0 {
		return &plugins.Result{
			Success: false,
			Service: "rabbitmq",
			Error:   fmt.Errorf("没有可用的测试凭据"),
		}
	}

	for _, cred := range credentials {
		// 检查上下文是否已取消
		select {
		case <-ctx.Done():
			return &plugins.Result{
				Success: false,
				Service: "rabbitmq",
				Error:   ctx.Err(),
			}
		default:
		}

		if p.testCredential(ctx, info, cred) {
			common.LogSuccess(fmt.Sprintf("RabbitMQ %s %s:%s", target, cred.Username, cred.Password))
			return &plugins.Result{
				Success:  true,
				Service:  "rabbitmq",
				Username: cred.Username,
				Password: cred.Password,
			}
		}
	}

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


// testAMQPProtocol 检测AMQP协议
func (p *RabbitMQPlugin) testAMQPProtocol(ctx context.Context, info *common.HostInfo) *plugins.Result {
	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
	
	// 连接到AMQP端口
	conn, err := net.DialTimeout("tcp", target, time.Duration(common.Timeout)*time.Second)
	if err != nil {
		return &plugins.Result{
			Success: false,
			Service: "rabbitmq",
			Error:   err,
		}
	}
	defer conn.Close()

	// 设置超时
	conn.SetDeadline(time.Now().Add(time.Duration(common.Timeout) * time.Second))

	// 发送AMQP协议头请求 (AMQP 0-9-1)
	amqpHeader := []byte{0x41, 0x4d, 0x51, 0x50, 0x00, 0x00, 0x09, 0x01}
	_, err = conn.Write(amqpHeader)
	if err != nil {
		return &plugins.Result{
			Success: false,
			Service: "rabbitmq",
			Error:   fmt.Errorf("发送AMQP握手失败: %v", err),
		}
	}

	// 读取服务器响应
	buffer := make([]byte, 32)
	n, err := conn.Read(buffer)
	if err != nil {
		return &plugins.Result{
			Success: false,
			Service: "rabbitmq",
			Error:   fmt.Errorf("读取AMQP响应失败: %v", err),
		}
	}

	// 调试信息（可选）
	common.LogDebug(fmt.Sprintf("RabbitMQ AMQP响应长度: %d, 前8字节: %v", n, buffer[:min(n, 8)]))

	// 检查AMQP协议头或连接开始帧
	if n >= 4 && string(buffer[:4]) == "AMQP" {
		// 服务器返回AMQP协议头
		banner := fmt.Sprintf("RabbitMQ AMQP %d.%d.%d", buffer[5], buffer[6], buffer[7])
		common.LogSuccess(fmt.Sprintf("RabbitMQ %s %s", target, banner))
		return &plugins.Result{
			Success: true,
			Service: "rabbitmq",
			Banner:  banner,
		}
	} else if n >= 8 && buffer[0] == 0x01 && buffer[7] == 0xCE {
		// Connection.Start方法帧 (AMQP 0-9-1)
		banner := "RabbitMQ AMQP 0-9-1"
		common.LogSuccess(fmt.Sprintf("RabbitMQ %s %s", target, banner))
		return &plugins.Result{
			Success: true,
			Service: "rabbitmq",
			Banner:  banner,
		}
	} else if n >= 8 && buffer[0] == 0x01 {
		// 可能是AMQP帧，但格式不同
		banner := "RabbitMQ AMQP"
		common.LogSuccess(fmt.Sprintf("RabbitMQ %s %s", target, banner))
		return &plugins.Result{
			Success: true,
			Service: "rabbitmq",
			Banner:  banner,
		}
	}

	return &plugins.Result{
		Success: false,
		Service: "rabbitmq",
		Error:   fmt.Errorf("非AMQP协议响应"),
	}
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func (p *RabbitMQPlugin) testCredential(ctx context.Context, info *common.HostInfo, cred plugins.Credential) bool {
	baseURL := fmt.Sprintf("http://%s:%s", info.Host, info.Ports)
	
	client := &http.Client{
		Timeout: time.Duration(common.Timeout) * time.Second,
	}

	req, err := http.NewRequestWithContext(ctx, "GET", baseURL+"/api/overview", nil)
	if err != nil {
		return false
	}

	req.SetBasicAuth(cred.Username, cred.Password)
	req.Header.Set("Content-Type", "application/json")

	resp, err := client.Do(req)
	if err != nil {
		return false
	}
	defer resp.Body.Close()

	return resp.StatusCode == 200
}








func (p *RabbitMQPlugin) identifyService(ctx context.Context, info *common.HostInfo) *plugins.Result {
	// 对于AMQP端口，检测AMQP协议
	if info.Ports == "5672" || info.Ports == "5671" {
		return p.testAMQPProtocol(ctx, info)
	}

	// 对于HTTP端口，检测管理界面
	return p.testManagementInterface(ctx, info)
}

// testManagementInterface 检测RabbitMQ管理界面
func (p *RabbitMQPlugin) testManagementInterface(ctx context.Context, info *common.HostInfo) *plugins.Result {
	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
	baseURL := fmt.Sprintf("http://%s:%s", info.Host, info.Ports)
	
	client := &http.Client{
		Timeout: time.Duration(common.Timeout) * time.Second,
	}

	req, err := http.NewRequestWithContext(ctx, "GET", baseURL, nil)
	if err != nil {
		return &plugins.Result{
			Success: false,
			Service: "rabbitmq",
			Error:   err,
		}
	}

	resp, err := client.Do(req)
	if err != nil {
		return &plugins.Result{
			Success: false,
			Service: "rabbitmq",
			Error:   err,
		}
	}
	defer resp.Body.Close()

	var banner string
	if resp.StatusCode == 200 || resp.StatusCode == 401 {
		body, _ := io.ReadAll(resp.Body)
		bodyStr := strings.ToLower(string(body))
		
		if strings.Contains(bodyStr, "rabbitmq") {
			banner = "RabbitMQ Management"
		} else if strings.Contains(bodyStr, "management") {
			banner = "RabbitMQ Management"
		} else {
			banner = "RabbitMQ Management"
		}
		common.LogSuccess(fmt.Sprintf("RabbitMQ %s %s", target, banner))
		return &plugins.Result{
			Success: true,
			Service: "rabbitmq",
			Banner:  banner,
		}
	} else {
		return &plugins.Result{
			Success: false,
			Service: "rabbitmq",
			Error:   fmt.Errorf("无法识别为RabbitMQ服务"),
		}
	}
}

func init() {
	plugins.RegisterWithPorts("rabbitmq", func() plugins.Plugin {
		return NewRabbitMQPlugin()
	}, []int{5672, 15672, 5671})
}