package services

import (
	"context"
	"crypto/des"
	"encoding/binary"
	"fmt"
	"net"
	"strings"
	"time"

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

type VNCPlugin struct {
	plugins.BasePlugin
}

func NewVNCPlugin() *VNCPlugin {
	return &VNCPlugin{
		BasePlugin: plugins.NewBasePlugin("vnc"),
	}
}



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

	if result := p.testUnauthAccess(ctx, info); result != nil && result.Success {
		common.LogSuccess(fmt.Sprintf("VNC %s 未授权访问", target))
		return result
	}

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

	for _, cred := range credentials {
		if p.testCredential(ctx, info, cred) {
			common.LogSuccess(fmt.Sprintf("VNC %s %s", target, cred.Password))

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

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


func (p *VNCPlugin) testUnauthAccess(ctx context.Context, info *common.HostInfo) *ScanResult {
	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
	
	// 检查发包限制
	if canSend, reason := common.CanSendPacket(); !canSend {
		common.LogError(fmt.Sprintf("VNC未授权检测 %s 受限: %s", target, reason))
		return nil
	}
	
	conn, err := net.DialTimeout("tcp", target, time.Duration(common.Timeout)*time.Second)
	if err != nil {
		common.IncrementTCPFailedPacketCount()
		return nil
	}
	common.IncrementTCPSuccessPacketCount()
	defer conn.Close()

	conn.SetDeadline(time.Now().Add(time.Duration(common.Timeout) * time.Second))

	_, err = p.performHandshake(conn)
	if err != nil {
		return nil
	}

	authTypes, err := p.readAuthTypes(conn)
	if err != nil {
		return nil
	}

	for _, authType := range authTypes {
		if authType == 1 {
			return &ScanResult{
				Success: true,
				Service: "vnc",
				Banner:  "未授权访问",
			}
		}
	}

	return nil
}


func (p *VNCPlugin) testCredential(ctx context.Context, info *common.HostInfo, cred Credential) bool {
	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
	
	// 检查发包限制
	if canSend, reason := common.CanSendPacket(); !canSend {
		common.LogError(fmt.Sprintf("VNC认证测试 %s 受限: %s", target, reason))
		return false
	}
	
	conn, err := net.DialTimeout("tcp", target, time.Duration(common.Timeout)*time.Second)
	if err != nil {
		common.IncrementTCPFailedPacketCount()
		return false
	}
	common.IncrementTCPSuccessPacketCount()
	defer conn.Close()

	conn.SetDeadline(time.Now().Add(time.Duration(common.Timeout) * time.Second))

	_, err = p.performHandshake(conn)
	if err != nil {
		return false
	}

	authTypes, err := p.readAuthTypes(conn)
	if err != nil {
		return false
	}

	var hasVNCAuth bool
	for _, authType := range authTypes {
		if authType == 2 {
			hasVNCAuth = true
			break
		}
	}

	if !hasVNCAuth {
		return false
	}

	if _, err := conn.Write([]byte{2}); err != nil {
		return false
	}

	challenge := make([]byte, 16)
	if _, err := conn.Read(challenge); err != nil {
		return false
	}

	key := make([]byte, 8)
	copy(key, []byte(cred.Password))

	response, err := p.encryptChallenge(challenge, key)
	if err != nil {
		return false
	}

	if _, err := conn.Write(response); err != nil {
		return false
	}

	result := make([]byte, 4)
	if _, err := conn.Read(result); err != nil {
		return false
	}

	return binary.BigEndian.Uint32(result) == 0
}


func (p *VNCPlugin) performHandshake(conn net.Conn) (string, error) {
	versionBuf := make([]byte, 12)
	if _, err := conn.Read(versionBuf); err != nil {
		return "", err
	}

	version := strings.TrimSpace(string(versionBuf))
	
	if _, err := conn.Write(versionBuf); err != nil {
		return "", err
	}

	return version, nil
}

func (p *VNCPlugin) readAuthTypes(conn net.Conn) ([]byte, error) {
	countBuf := make([]byte, 1)
	if _, err := conn.Read(countBuf); err != nil {
		return nil, err
	}

	count := countBuf[0]
	if count == 0 {
		return nil, fmt.Errorf("无可用认证类型")
	}

	authTypes := make([]byte, count)
	if _, err := conn.Read(authTypes); err != nil {
		return nil, err
	}

	return authTypes, nil
}

func (p *VNCPlugin) encryptChallenge(challenge, key []byte) ([]byte, error) {
	reversedKey := make([]byte, 8)
	for i := 0; i < 8; i++ {
		reversedKey[i] = p.reverseBits(key[i])
	}

	cipher, err := des.NewCipher(reversedKey)
	if err != nil {
		return nil, err
	}

	response := make([]byte, 16)
	cipher.Encrypt(response[0:8], challenge[0:8])
	cipher.Encrypt(response[8:16], challenge[8:16])

	return response, nil
}

func (p *VNCPlugin) reverseBits(b byte) byte {
	var result byte
	for i := 0; i < 8; i++ {
		result = (result << 1) | ((b >> i) & 1)
	}
	return result
}


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

	// 检查发包限制
	if canSend, reason := common.CanSendPacket(); !canSend {
		common.LogError(fmt.Sprintf("VNC识别 %s 受限: %s", target, reason))
		return &ScanResult{
			Success: false,
			Service: "vnc",
			Error:   fmt.Errorf("发包受限: %s", reason),
		}
	}

	conn, err := net.DialTimeout("tcp", target, time.Duration(common.Timeout)*time.Second)
	if err != nil {
		common.IncrementTCPFailedPacketCount()
		return &ScanResult{
			Success: false,
			Service: "vnc",
			Error:   err,
		}
	}
	common.IncrementTCPSuccessPacketCount()
	defer conn.Close()

	conn.SetDeadline(time.Now().Add(time.Duration(common.Timeout) * time.Second))

	version, err := p.performHandshake(conn)
	if err != nil {
		return &ScanResult{
			Success: false,
			Service: "vnc",
			Error:   err,
		}
	}

	banner := fmt.Sprintf("VNC远程桌面服务 (%s)", version)
	common.LogSuccess(fmt.Sprintf("VNC %s %s", target, banner))

	return &ScanResult{
		Success: true,
		Service: "vnc",
		Banner:  banner,
	}
}

func init() {
	// 使用高效注册方式：直接传递端口信息，避免实例创建
	RegisterPluginWithPorts("vnc", func() Plugin {
		return NewVNCPlugin()
	}, []int{5900, 5901, 5902, 5903, 5904, 5905, 5906, 5907, 5908, 5909})
}