package services

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

	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
	"github.com/shadow1ng/fscan/common"
	"github.com/shadow1ng/fscan/plugins"
)

type MongoDBPlugin struct {
	plugins.BasePlugin
}

func NewMongoDBPlugin() *MongoDBPlugin {
	return &MongoDBPlugin{
		BasePlugin: plugins.NewBasePlugin("mongodb"),
	}
}

func (p *MongoDBPlugin) 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)
	}

	// 首先检测未授权访问
	isUnauth, err := p.mongodbUnauth(ctx, info)
	if err != nil {
		return &plugins.Result{
			Success: false,
			Service: "mongodb",
			Error:   err,
		}
	}

	if isUnauth {
		common.LogSuccess(fmt.Sprintf("MongoDB %s 未授权访问", target))
		return &plugins.Result{
			Success: true,
			Service: "mongodb",
			Banner:  "未授权访问",
		}
	}

	// 如果需要认证，尝试常见凭据
	credentials := plugins.GenerateCredentials("mongodb")
	
	for _, cred := range credentials {
		if p.testMongoCredential(ctx, info, cred) {
			common.LogSuccess(fmt.Sprintf("MongoDB %s %s:%s", target, cred.Username, cred.Password))
			return &plugins.Result{
				Success:  true,
				Service:  "mongodb",
				Username: cred.Username,
				Password: cred.Password,
			}
		}
	}

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

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

	// 尝试检测MongoDB服务
	isUnauth, err := p.mongodbUnauth(ctx, info)
	if err != nil {
		return &plugins.Result{
			Success: false,
			Service: "mongodb",
			Error:   err,
		}
	}

	// 如果能获得MongoDB响应（无论是否授权），都说明是MongoDB服务
	if isUnauth {
		common.LogSuccess(fmt.Sprintf("MongoDB %s 未授权访问", target))
		return &plugins.Result{
			Success: true,
			Service: "mongodb",
			Banner:  "未授权访问",
		}
	} else {
		common.LogSuccess(fmt.Sprintf("MongoDB %s 需要认证", target))
		return &plugins.Result{
			Success: true,
			Service: "mongodb",
			Banner:  "需要认证",
		}
	}
}

// mongodbUnauth 检测MongoDB未授权访问 - 基于原始工作版本
func (p *MongoDBPlugin) mongodbUnauth(ctx context.Context, info *common.HostInfo) (bool, error) {
	msgPacket := p.createOpMsgPacket()
	queryPacket := p.createOpQueryPacket()

	realhost := fmt.Sprintf("%s:%s", info.Host, info.Ports)
	
	// 尝试OP_MSG查询
	reply, err := p.checkMongoAuth(ctx, realhost, msgPacket)
	if err != nil {
		// 失败则尝试OP_QUERY查询
		reply, err = p.checkMongoAuth(ctx, realhost, queryPacket)
		if err != nil {
			return false, err
		}
	}

	// 检查响应结果 - 基于原始版本的检测逻辑
	if strings.Contains(reply, "totalLinesWritten") {
		return true, nil
	}

	// 如果能收到响应但不包含预期内容，说明是MongoDB但需要认证
	if len(reply) > 0 {
		return false, nil // 是MongoDB，但需要认证
	}

	return false, fmt.Errorf("无法识别为MongoDB服务")
}

// checkMongoAuth 检查MongoDB认证状态 - 基于原始工作版本
func (p *MongoDBPlugin) checkMongoAuth(ctx context.Context, address string, packet []byte) (string, error) {
	// 创建连接超时上下文
	connCtx, cancel := context.WithTimeout(ctx, time.Duration(common.Timeout)*time.Second)
	defer cancel()

	// 使用带超时的连接
	var d net.Dialer
	conn, err := d.DialContext(connCtx, "tcp", address)
	if err != nil {
		return "", fmt.Errorf("连接失败: %v", err)
	}
	defer conn.Close()

	// 检查上下文是否已取消
	select {
	case <-ctx.Done():
		return "", ctx.Err()
	default:
	}

	// 设置读写超时
	if err := conn.SetDeadline(time.Now().Add(time.Duration(common.Timeout) * time.Second)); err != nil {
		return "", fmt.Errorf("设置超时失败: %v", err)
	}

	// 发送查询包
	if _, err := conn.Write(packet); err != nil {
		return "", fmt.Errorf("发送查询失败: %v", err)
	}

	// 再次检查上下文是否已取消
	select {
	case <-ctx.Done():
		return "", ctx.Err()
	default:
	}

	// 读取响应
	reply := make([]byte, 2048)
	count, err := conn.Read(reply)
	if err != nil && err != io.EOF {
		return "", fmt.Errorf("读取响应失败: %v", err)
	}

	if count == 0 {
		return "", fmt.Errorf("收到空响应")
	}

	return string(reply[:count]), nil
}

// createOpMsgPacket 创建OP_MSG查询包 - 直接使用原始工作版本
func (p *MongoDBPlugin) createOpMsgPacket() []byte {
	return []byte{
		0x69, 0x00, 0x00, 0x00, // messageLength
		0x39, 0x00, 0x00, 0x00, // requestID
		0x00, 0x00, 0x00, 0x00, // responseTo
		0xdd, 0x07, 0x00, 0x00, // opCode OP_MSG
		0x00, 0x00, 0x00, 0x00, // flagBits
		// sections db.adminCommand({getLog: "startupWarnings"})
		0x00, 0x54, 0x00, 0x00, 0x00, 0x02, 0x67, 0x65, 0x74, 0x4c, 0x6f, 0x67, 0x00, 0x10, 0x00, 0x00, 0x00, 0x73, 0x74, 0x61, 0x72, 0x74, 0x75, 0x70, 0x57, 0x61, 0x72, 0x6e, 0x69, 0x6e, 0x67, 0x73, 0x00, 0x02, 0x24, 0x64, 0x62, 0x00, 0x06, 0x00, 0x00, 0x00, 0x61, 0x64, 0x6d, 0x69, 0x6e, 0x00, 0x03, 0x6c, 0x73, 0x69, 0x64, 0x00, 0x1e, 0x00, 0x00, 0x00, 0x05, 0x69, 0x64, 0x00, 0x10, 0x00, 0x00, 0x00, 0x04, 0x6e, 0x81, 0xf8, 0x8e, 0x37, 0x7b, 0x4c, 0x97, 0x84, 0x4e, 0x90, 0x62, 0x5a, 0x54, 0x3c, 0x93, 0x00, 0x00,
	}
}

// createOpQueryPacket 创建OP_QUERY查询包 - 直接使用原始工作版本
func (p *MongoDBPlugin) createOpQueryPacket() []byte {
	return []byte{
		0x48, 0x00, 0x00, 0x00, // messageLength
		0x02, 0x00, 0x00, 0x00, // requestID
		0x00, 0x00, 0x00, 0x00, // responseTo
		0xd4, 0x07, 0x00, 0x00, // opCode OP_QUERY
		0x00, 0x00, 0x00, 0x00, // flags
		0x61, 0x64, 0x6d, 0x69, 0x6e, 0x2e, 0x24, 0x63, 0x6d, 0x64, 0x00, // fullCollectionName admin.$cmd
		0x00, 0x00, 0x00, 0x00, // numberToSkip
		0x01, 0x00, 0x00, 0x00, // numberToReturn
		// query db.adminCommand({getLog: "startupWarnings"})
		0x21, 0x00, 0x00, 0x00, 0x2, 0x67, 0x65, 0x74, 0x4c, 0x6f, 0x67, 0x00, 0x10, 0x00, 0x00, 0x00, 0x73, 0x74, 0x61, 0x72, 0x74, 0x75, 0x70, 0x57, 0x61, 0x72, 0x6e, 0x69, 0x6e, 0x67, 0x73, 0x00, 0x00,
	}
}

// testMongoCredential 使用官方MongoDB驱动测试凭据
func (p *MongoDBPlugin) testMongoCredential(ctx context.Context, info *common.HostInfo, cred plugins.Credential) bool {
	// 构建MongoDB连接URI
	var uri string
	if cred.Username != "" && cred.Password != "" {
		uri = fmt.Sprintf("mongodb://%s:%s@%s:%s/?connectTimeoutMS=%d000&serverSelectionTimeoutMS=%d000",
			cred.Username, cred.Password, info.Host, info.Ports, common.Timeout, common.Timeout)
	} else if cred.Username != "" {
		// 对于有用户名但密码为空的情况，仍然尝试认证
		uri = fmt.Sprintf("mongodb://%s:@%s:%s/?connectTimeoutMS=%d000&serverSelectionTimeoutMS=%d000",
			cred.Username, info.Host, info.Ports, common.Timeout, common.Timeout)
	} else {
		// 无用户名的情况，尝试无认证连接
		uri = fmt.Sprintf("mongodb://%s:%s/?connectTimeoutMS=%d000&serverSelectionTimeoutMS=%d000",
			info.Host, info.Ports, common.Timeout, common.Timeout)
	}


	// 创建客户端选项
	clientOptions := options.Client().ApplyURI(uri)
	
	// 创建带超时的上下文
	authCtx, cancel := context.WithTimeout(ctx, time.Duration(common.Timeout)*time.Second)
	defer cancel()

	// 连接到MongoDB
	client, err := mongo.Connect(authCtx, clientOptions)
	if err != nil {
		return false
	}
	defer client.Disconnect(authCtx)

	// 测试连接 - 尝试ping数据库
	err = client.Ping(authCtx, nil)
	if err != nil {
		return false
	}

	return true
}

func init() {
	plugins.RegisterWithPorts("mongodb", func() plugins.Plugin {
		return NewMongoDBPlugin()
	}, []int{27017, 27018, 27019})
}