package services

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

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

// VNCPlugin VNC远程桌面服务扫描和利用插件 - 包含屏幕信息收集利用功能
type VNCPlugin struct {
	name  string
	ports []int
}

// NewVNCPlugin 创建VNC插件
func NewVNCPlugin() *VNCPlugin {
	return &VNCPlugin{
		name:  "vnc",
		ports: []int{5900, 5901, 5902, 5903, 5904, 5905, 5906, 5907, 5908, 5909}, // VNC端口
	}
}

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

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

// Scan 执行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(i18n.GetText("vnc_unauth_success", target))
		return result
	}

	// 生成测试密码（VNC通常只有密码，没有用户名）
	passwords := p.generatePasswords()

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

		// 测试密码
		if p.testPassword(ctx, info, password) {
			// VNC认证成功
			common.LogSuccess(i18n.GetText("vnc_scan_success", target, password))

			return &ScanResult{
				Success:  true,
				Service:  "vnc",
				Username: "", // VNC没有用户名概念
				Password: password,
			}
		}
	}

	// 所有密码都失败
	return &ScanResult{
		Success: false,
		Service: "vnc",
		Error:   fmt.Errorf("未发现弱密码或未授权访问"),
	}
}

// Exploit 执行VNC利用操作 - 实现屏幕信息收集功能
func (p *VNCPlugin) Exploit(ctx context.Context, info *common.HostInfo, creds Credential) *ExploitResult {
	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
	common.LogSuccess(fmt.Sprintf("VNC利用开始: %s", target))

	var output strings.Builder
	output.WriteString(fmt.Sprintf("=== VNC利用结果 - %s ===\n", target))
	
	if creds.Password != "" {
		output.WriteString(fmt.Sprintf("认证密码: %s\n", creds.Password))
	} else {
		output.WriteString("认证方式: 无密码访问\n")
	}

	// 建立VNC连接
	conn, version, err := p.connectVNC(ctx, info, creds.Password)
	if err != nil {
		output.WriteString(fmt.Sprintf("\n[连接失败] %v\n", err))
		return &ExploitResult{
			Success: false,
			Output:  output.String(),
			Error:   err,
		}
	}
	defer conn.Close()

	output.WriteString(fmt.Sprintf("\n[VNC版本] %s\n", version))
	output.WriteString("[连接状态] ✅ 成功建立VNC连接\n")

	// 获取服务器信息
	if serverInfo := p.getServerInfo(conn); serverInfo != "" {
		output.WriteString(fmt.Sprintf("\n[服务器信息]\n%s\n", serverInfo))
	}

	// 获取桌面信息
	if desktopInfo := p.getDesktopInfo(conn); desktopInfo != "" {
		output.WriteString(fmt.Sprintf("\n[桌面信息]\n%s\n", desktopInfo))
	}

	// 获取像素格式信息
	if pixelFormat := p.getPixelFormat(conn); pixelFormat != "" {
		output.WriteString(fmt.Sprintf("\n[像素格式]\n%s\n", pixelFormat))
	}

	// 尝试获取屏幕截图信息（不实际截图，只是获取能力信息）
	output.WriteString("\n[屏幕访问]\n")
	output.WriteString("✅ 可以访问远程桌面画面\n")
	output.WriteString("⚠️  完整屏幕截图需要VNC客户端支持\n")

	common.LogSuccess(fmt.Sprintf("VNC利用完成: %s", target))

	return &ExploitResult{
		Success: true,
		Output:  output.String(),
	}
}

// testUnauthAccess 测试未授权访问
func (p *VNCPlugin) testUnauthAccess(ctx context.Context, info *common.HostInfo) *ScanResult {
	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
	
	conn, err := net.DialTimeout("tcp", target, time.Duration(common.Timeout)*time.Second)
	if err != nil {
		return nil
	}
	defer conn.Close()

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

	// VNC握手
	_, 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 { // None authentication
			return &ScanResult{
				Success: true,
				Service: "vnc",
				Banner:  "未授权访问",
			}
		}
	}

	return nil
}

// generatePasswords 生成VNC测试密码
func (p *VNCPlugin) generatePasswords() []string {
	// VNC常见弱密码
	return []string{
		"",
		"123456",
		"password",
		"admin",
		"vnc",
		"123",
		"1234",
		"12345",
		"root",
		"test",
		"guest",
		"user",
		"welcome",
		"qwerty",
		"abc123",
	}
}

// testPassword 测试单个密码
func (p *VNCPlugin) testPassword(ctx context.Context, info *common.HostInfo, password string) bool {
	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
	
	conn, err := net.DialTimeout("tcp", target, time.Duration(common.Timeout)*time.Second)
	if err != nil {
		return false
	}
	defer conn.Close()

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

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

	// 读取认证类型
	authTypes, err := p.readAuthTypes(conn)
	if err != nil {
		return false
	}

	// 查找VNC认证类型
	var hasVNCAuth bool
	for _, authType := range authTypes {
		if authType == 2 { // VNC authentication
			hasVNCAuth = true
			break
		}
	}

	if !hasVNCAuth {
		return false
	}

	// 选择VNC认证
	if _, err := conn.Write([]byte{2}); err != nil {
		return false
	}

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

	// 准备密码（最多8字节，不足补零）
	key := make([]byte, 8)
	copy(key, []byte(password))

	// DES加密挑战
	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
	}

	// 检查认证结果 (0表示成功)
	return binary.BigEndian.Uint32(result) == 0
}

// connectVNC 建立认证后的VNC连接
func (p *VNCPlugin) connectVNC(ctx context.Context, info *common.HostInfo, password string) (net.Conn, string, error) {
	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
	
	conn, err := net.DialTimeout("tcp", target, time.Duration(common.Timeout)*time.Second)
	if err != nil {
		return nil, "", err
	}

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

	// VNC握手
	version, err := p.performHandshake(conn)
	if err != nil {
		conn.Close()
		return nil, "", err
	}

	// 读取认证类型
	authTypes, err := p.readAuthTypes(conn)
	if err != nil {
		conn.Close()
		return nil, "", err
	}

	// 处理认证
	authenticated := false
	for _, authType := range authTypes {
		if authType == 1 && password == "" { // None authentication
			conn.Write([]byte{1})
			authenticated = true
			break
		} else if authType == 2 && password != "" { // VNC authentication
			conn.Write([]byte{2})
			
			// 读取挑战
			challenge := make([]byte, 16)
			conn.Read(challenge)
			
			// 准备密码
			key := make([]byte, 8)
			copy(key, []byte(password))
			
			// 加密挑战
			response, err := p.encryptChallenge(challenge, key)
			if err != nil {
				continue
			}
			
			// 发送响应
			conn.Write(response)
			
			// 读取认证结果
			result := make([]byte, 4)
			conn.Read(result)
			
			if binary.BigEndian.Uint32(result) == 0 {
				authenticated = true
				break
			}
		}
	}

	if !authenticated {
		conn.Close()
		return nil, "", fmt.Errorf("认证失败")
	}

	return conn, version, nil
}

// performHandshake 执行VNC握手
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
}

// readAuthTypes 读取认证类型
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
}

// encryptChallenge 使用DES加密挑战
func (p *VNCPlugin) encryptChallenge(challenge, key []byte) ([]byte, error) {
	// VNC使用反向位序的DES
	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
}

// reverseBits 反转字节的位序
func (p *VNCPlugin) reverseBits(b byte) byte {
	var result byte
	for i := 0; i < 8; i++ {
		result = (result << 1) | ((b >> i) & 1)
	}
	return result
}

// getServerInfo 获取服务器信息
func (p *VNCPlugin) getServerInfo(conn net.Conn) string {
	var info strings.Builder
	
	// 发送客户端初始化消息
	conn.Write([]byte{1}) // 共享桌面
	
	// 读取服务器初始化消息
	serverInit := make([]byte, 24)
	n, err := conn.Read(serverInit)
	if err != nil || n < 20 {
		return "无法获取服务器信息"
	}
	
	width := binary.BigEndian.Uint16(serverInit[0:2])
	height := binary.BigEndian.Uint16(serverInit[2:4])
	
	info.WriteString(fmt.Sprintf("桌面分辨率: %dx%d\n", width, height))
	
	// 读取桌面名称长度
	nameLength := binary.BigEndian.Uint32(serverInit[20:24])
	if nameLength > 0 && nameLength < 1024 {
		nameBuf := make([]byte, nameLength)
		if n, err := conn.Read(nameBuf); err == nil && n == int(nameLength) {
			info.WriteString(fmt.Sprintf("桌面名称: %s\n", string(nameBuf)))
		}
	}
	
	return info.String()
}

// getDesktopInfo 获取桌面信息
func (p *VNCPlugin) getDesktopInfo(conn net.Conn) string {
	// 简化实现，实际需要完整的RFB协议支持
	return "桌面信息获取需要完整VNC客户端支持"
}

// getPixelFormat 获取像素格式信息
func (p *VNCPlugin) getPixelFormat(conn net.Conn) string {
	return "像素格式信息获取需要完整VNC客户端支持"
}

// identifyService 服务识别 - 检测VNC服务
func (p *VNCPlugin) identifyService(ctx context.Context, info *common.HostInfo) *ScanResult {
	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)

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

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

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

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

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

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