package vnc

import (
	"context"
	"fmt"
	"time"

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

// VNCConnector VNC连接器
type VNCConnector struct {
	host string
	port string
}

// NewVNCConnector 创建VNC连接器
func NewVNCConnector() *VNCConnector {
	return &VNCConnector{}
}

// Connect 连接到VNC服务
func (c *VNCConnector) Connect(ctx context.Context, info *common.HostInfo) (interface{}, error) {
	c.host = info.Host
	c.port = info.Ports
	
	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
	
	// 尝试连接到VNC服务
	conn, err := common.WrapperTcpWithTimeout("tcp", target, time.Duration(common.Timeout)*time.Second)
	if err != nil {
		return nil, fmt.Errorf("VNC连接失败: %v", err)
	}
	defer conn.Close()
	
	// 读取VNC协议版本
	versionBuf := make([]byte, 12)
	_, err = conn.Read(versionBuf)
	if err != nil {
		return nil, fmt.Errorf("读取VNC版本失败: %v", err)
	}
	
	versionStr := string(versionBuf)
	if !c.isValidVNCVersion(versionStr) {
		return nil, fmt.Errorf("无效的VNC协议版本: %s", versionStr)
	}
	
	return map[string]interface{}{
		"version": versionStr,
		"target":  target,
	}, nil
}

// Authenticate VNC认证
func (c *VNCConnector) Authenticate(ctx context.Context, conn interface{}, cred *base.Credential) error {
	connectionInfo, ok := conn.(map[string]interface{})
	if !ok {
		return fmt.Errorf("无效的连接信息")
	}
	
	target := connectionInfo["target"].(string)
	
	// 重新连接进行认证测试
	tcpConn, err := common.WrapperTcpWithTimeout("tcp", target, time.Duration(common.Timeout)*time.Second)
	if err != nil {
		return fmt.Errorf("重连VNC失败: %v", err)
	}
	defer tcpConn.Close()
	
	// 读取并回应版本
	versionBuf := make([]byte, 12)
	_, err = tcpConn.Read(versionBuf)
	if err != nil {
		return fmt.Errorf("读取版本失败: %v", err)
	}
	
	// 发送版本响应
	_, err = tcpConn.Write(versionBuf)
	if err != nil {
		return fmt.Errorf("发送版本响应失败: %v", err)
	}
	
	// 读取安全类型
	secTypeBuf := make([]byte, 1)
	_, err = tcpConn.Read(secTypeBuf)
	if err != nil {
		return fmt.Errorf("读取安全类型失败: %v", err)
	}
	
	secTypeCount := int(secTypeBuf[0])
	if secTypeCount == 0 {
		return fmt.Errorf("VNC服务器拒绝连接")
	}
	
	// 读取安全类型列表
	secTypes := make([]byte, secTypeCount)
	_, err = tcpConn.Read(secTypes)
	if err != nil {
		return fmt.Errorf("读取安全类型列表失败: %v", err)
	}
	
	// 检查是否支持无认证 (类型1) 或VNC认证 (类型2)
	var selectedSecType byte = 0
	for _, secType := range secTypes {
		if secType == 1 { // 无认证
			selectedSecType = 1
			break
		} else if secType == 2 { // VNC认证
			selectedSecType = 2
			break
		}
	}
	
	if selectedSecType == 0 {
		return fmt.Errorf("不支持的VNC安全类型")
	}
	
	// 发送选择的安全类型
	_, err = tcpConn.Write([]byte{selectedSecType})
	if err != nil {
		return fmt.Errorf("发送安全类型失败: %v", err)
	}
	
	if selectedSecType == 1 {
		// 无认证模式，直接成功
		common.LogSuccess(fmt.Sprintf("VNC无认证访问成功: %s", target))
		return nil
	} else if selectedSecType == 2 {
		// VNC认证模式，需要密码
		if cred.Password == "" {
			return fmt.Errorf("VNC认证需要密码")
		}
		
		// 读取质询
		challenge := make([]byte, 16)
		_, err = tcpConn.Read(challenge)
		if err != nil {
			return fmt.Errorf("读取VNC质询失败: %v", err)
		}
		
		// 这里应该实现DES加密，但为了简化，我们只检测是否需要密码
		// 实际实现中需要使用DES加密算法对质询进行加密
		return fmt.Errorf("VNC密码认证需要进一步实现")
	}
	
	return fmt.Errorf("未知的安全类型: %d", selectedSecType)
}

// Close 关闭连接
func (c *VNCConnector) Close(conn interface{}) error {
	// VNC连接器不需要保持长连接
	return nil
}

// GetDefaultCredentials 获取VNC默认凭据
func (c *VNCConnector) GetDefaultCredentials() []*base.Credential {
	return []*base.Credential{
		{Username: "", Password: ""},        // 无密码
		{Username: "", Password: "admin"},   
		{Username: "", Password: "password"},
		{Username: "", Password: "123456"},
		{Username: "", Password: "vnc"},
		{Username: "", Password: "root"},
	}
}

// isValidVNCVersion 检查是否为有效的VNC版本
func (c *VNCConnector) isValidVNCVersion(version string) bool {
	// VNC版本格式: RFB xxx.yyy\n
	if len(version) < 12 {
		return false
	}
	
	return version[:4] == "RFB " && (version[11] == '\n' || version[11] == '\r')
}

// GetSupportedVersions 获取支持的VNC版本
func (c *VNCConnector) GetSupportedVersions() []string {
	return []string{"RFB 003.003", "RFB 003.007", "RFB 003.008"}
}