//go:build windows

package remote

import (
	"bytes"
	"crypto/rand"
	"golang.org/x/sys/windows/registry"
	"math/big"
	"os"
)

// GetRemoteInfo 一次性获取远程控制相关信息（服务状态/端口/密码/主机信息）
func GetRemoteInfo() Info {
	info := Info{}

	status, _ := GetVncServiceStatus()
	info.VncState = int(status)

	status, _ = GetRdpServiceStatus()
	info.RdpState = int(status)

	port, _ := GetRdpPort()
	info.RdpPort = port

	port, _ = GetVncPort()
	info.VncPort = port

	password, _ := GetVncPassword()
	info.VncPassword = password

	hostname, _ := os.Hostname()
	info.Hostname = hostname
	info.IpAddr = getIPs()
	info.MacAddr = getMacAddrs()

	return info
}

// GetVncPort 从注册表读取 TightVNC 主端口
func GetVncPort() (int, error) {
	var key registry.Key
	var err error
	key, err = registry.OpenKey(registry.LOCAL_MACHINE, `SOFTWARE\TightVNC\Server`, registry.ALL_ACCESS)
	if err != nil {
		return 0, err
	}
	var val uint64
	val, _, err = key.GetIntegerValue("RfbPort")
	if err != nil {
		return 0, err
	}
	var port int = int(val)
	return port, nil
}

// GetRdpPort 从注册表读取远程桌面端口
func GetRdpPort() (int, error) {
	var key registry.Key
	var err error
	key, err = registry.OpenKey(registry.LOCAL_MACHINE, `SYSTEM\CurrentControlSet\Control\Terminal Server\WinStations\RDP-Tcp`, registry.ALL_ACCESS)
	if err != nil {
		return 0, err
	}
	var val uint64
	val, _, err = key.GetIntegerValue("PortNumber")
	if err != nil {
		return 0, err
	}
	var port int = int(val)
	return port, nil
}

// GetVncPassword 读取并解密 TightVNC 连接密码
func GetVncPassword() (string, error) {
	var key registry.Key
	var err error
	key, err = registry.OpenKey(registry.LOCAL_MACHINE, `SOFTWARE\TightVNC\Server`, registry.ALL_ACCESS)
	if err != nil {
		return "", err
	}

	val, _, err := key.GetBinaryValue("Password")
	if err != nil {
		return "", err
	}
	if len(val) == 0 {
		return "", err
	}
	password := VncPasswordDecrypt(val)
	if err != nil {
		return "", err
	}
	return password, nil
}

// VncServerDefaultConfig 生成 TightVNC 服务端默认配置（连接 RDP 会话）
// randPass 为 true 时生成 8 位随机密码，否则使用固定密码
func VncServerDefaultConfig(randPass bool) VncServerConfig {
	password := "tvncpass"
	if randPass {
		password = createRandomString(8)
	}
	config := VncServerConfig{
		ExtraPorts:                  "",
		QueryTimeout:                30,
		QueryAcceptOnTimeout:        0,
		LocalInputPriorityTimeout:   3,
		LocalInputPriority:          1,
		BlockRemoteInput:            0,
		BlockLocalInput:             0,
		IpAccessControl:             "",
		RfbPort:                     5900,
		HttpPort:                    5800,
		DisconnectAction:            0,
		AcceptRfbConnections:        1,
		UseVncAuthentication:        1,
		UseControlAuthentication:    0,
		RepeatControlAuthentication: 0,
		LoopbackOnly:                0,
		AcceptHttpConnections:       0,
		LogLevel:                    0,
		EnableFileTransfers:         1,
		RemoveWallpaper:             1,
		UseD3D:                      1,
		UseMirrorDriver:             1,
		EnableUrlParams:             1,
		Password:                    VncPasswordCrypt(password),
		AlwaysShared:                0,
		NeverShared:                 0,
		DisconnectClients:           1,
		PollingInterval:             1000,
		AllowLoopback:               1,
		VideoRecognitionInterval:    3000,
		GrabTransparentWindows:      1,
		SaveLogToAllUsersPath:       0,
		RunControlInterface:         0,
		IdleTimeout:                 0,
		VideoClasses:                "",
		VideoRects:                  "",
		ConnectToRdp:                1,
	}

	return config
}

// VncClientDefaultConfig 生成 TightVNC 客户端默认配置（不连接 RDP 会话）
// randPass 为 true 时生成 8 位随机密码，否则使用固定密码
func VncClientDefaultConfig(randPass bool) VncServerConfig {
	password := "tvncpass"
	if randPass {
		password = createRandomString(8)
	}
	config := VncServerConfig{
		ExtraPorts:                  "",
		QueryTimeout:                30,
		QueryAcceptOnTimeout:        0,
		LocalInputPriorityTimeout:   3,
		LocalInputPriority:          1,
		BlockRemoteInput:            0,
		BlockLocalInput:             0,
		IpAccessControl:             "",
		RfbPort:                     5900,
		HttpPort:                    5800,
		DisconnectAction:            0,
		AcceptRfbConnections:        1,
		UseVncAuthentication:        1,
		UseControlAuthentication:    0,
		RepeatControlAuthentication: 0,
		LoopbackOnly:                0,
		AcceptHttpConnections:       0,
		LogLevel:                    0,
		EnableFileTransfers:         1,
		RemoveWallpaper:             0,
		UseD3D:                      1,
		UseMirrorDriver:             1,
		EnableUrlParams:             1,
		Password:                    VncPasswordCrypt(password),
		AlwaysShared:                0,
		NeverShared:                 0,
		DisconnectClients:           1,
		PollingInterval:             1000,
		AllowLoopback:               1,
		VideoRecognitionInterval:    3000,
		GrabTransparentWindows:      1,
		SaveLogToAllUsersPath:       0,
		RunControlInterface:         0,
		IdleTimeout:                 0,
		VideoClasses:                "",
		VideoRects:                  "",
		ConnectToRdp:                0,
	}

	return config
}

// WriteVncServerConfig 将 TightVNC 服务端配置写入注册表（项不存在时自动创建）
func WriteVncServerConfig(config VncServerConfig) error {

	key, err := registry.OpenKey(registry.LOCAL_MACHINE, "SOFTWARE\\TightVNC\\Server", registry.ALL_ACCESS)
	if err != nil {
		if err == registry.ErrNotExist {
			// 项不存在 创建它
			key, _, err = registry.CreateKey(registry.LOCAL_MACHINE, "SOFTWARE\\TightVNC\\Server", registry.ALL_ACCESS)
			if err != nil {
				return err
			}
		} else {
			return err
		}
	}

	defer key.Close()

	err = key.SetStringValue("ExtraPorts", config.ExtraPorts)
	if err != nil {
		return err
	}
	err = key.SetDWordValue("QueryTimeout", config.QueryTimeout)
	if err != nil {
		return err
	}
	err = key.SetDWordValue("QueryAcceptOnTimeout", config.QueryAcceptOnTimeout)
	if err != nil {
		return err
	}
	err = key.SetDWordValue("LocalInputPriorityTimeout", config.LocalInputPriorityTimeout)
	if err != nil {
		return err
	}
	err = key.SetDWordValue("LocalInputPriority", config.LocalInputPriority)
	if err != nil {
		return err
	}
	err = key.SetDWordValue("BlockRemoteInput", config.BlockRemoteInput)
	if err != nil {
		return err
	}
	err = key.SetDWordValue("BlockLocalInput", config.BlockLocalInput)
	if err != nil {
		return err
	}
	err = key.SetStringValue("IpAccessControl", config.IpAccessControl)
	if err != nil {
		return err
	}
	err = key.SetDWordValue("RfbPort", config.RfbPort)
	if err != nil {
		return err
	}
	err = key.SetDWordValue("HttpPort", config.HttpPort)
	if err != nil {
		return err
	}
	err = key.SetDWordValue("DisconnectAction", config.DisconnectAction)
	if err != nil {
		return err
	}
	err = key.SetDWordValue("AcceptRfbConnections", config.AcceptRfbConnections)
	if err != nil {
		return err
	}
	err = key.SetDWordValue("UseVncAuthentication", config.UseVncAuthentication)
	if err != nil {
		return err
	}
	err = key.SetDWordValue("UseControlAuthentication", config.UseControlAuthentication)
	if err != nil {
		return err
	}
	err = key.SetDWordValue("RepeatControlAuthentication", config.RepeatControlAuthentication)
	if err != nil {
		return err
	}
	err = key.SetDWordValue("LoopbackOnly", config.LoopbackOnly)
	if err != nil {
		return err
	}
	err = key.SetDWordValue("AcceptHttpConnections", config.AcceptHttpConnections)
	if err != nil {
		return err
	}
	err = key.SetDWordValue("LogLevel", config.LogLevel)
	if err != nil {
		return err
	}
	err = key.SetDWordValue("EnableFileTransfers", config.EnableFileTransfers)
	if err != nil {
		return err
	}
	err = key.SetDWordValue("RemoveWallpaper", config.RemoveWallpaper)
	if err != nil {
		return err
	}
	err = key.SetDWordValue("UseD3D", config.UseD3D)
	if err != nil {
		return err
	}
	err = key.SetDWordValue("UseMirrorDriver", config.UseMirrorDriver)
	if err != nil {
		return err
	}
	err = key.SetDWordValue("EnableUrlParams", config.EnableUrlParams)
	if err != nil {
		return err
	}
	err = key.SetBinaryValue("Password", config.Password)
	if err != nil {
		return err
	}
	err = key.SetDWordValue("AlwaysShared", config.AlwaysShared)
	if err != nil {
		return err
	}
	err = key.SetDWordValue("NeverShared", config.NeverShared)
	if err != nil {
		return err
	}
	err = key.SetDWordValue("DisconnectClients", config.DisconnectClients)
	if err != nil {
		return err
	}
	err = key.SetDWordValue("PollingInterval", config.PollingInterval)
	if err != nil {
		return err
	}
	err = key.SetDWordValue("AllowLoopback", config.AllowLoopback)
	if err != nil {
		return err
	}
	err = key.SetDWordValue("VideoRecognitionInterval", config.VideoRecognitionInterval)
	if err != nil {
		return err
	}
	err = key.SetDWordValue("GrabTransparentWindows", config.GrabTransparentWindows)
	if err != nil {
		return err
	}
	err = key.SetDWordValue("SaveLogToAllUsersPath", config.SaveLogToAllUsersPath)
	if err != nil {
		return err
	}
	err = key.SetDWordValue("RunControlInterface", config.RunControlInterface)
	if err != nil {
		return err
	}
	err = key.SetDWordValue("IdleTimeout", config.IdleTimeout)
	if err != nil {
		return err
	}
	err = key.SetStringValue("VideoClasses", config.VideoClasses)
	if err != nil {
		return err
	}
	err = key.SetStringValue("VideoRects", config.VideoRects)
	if err != nil {
		return err
	}
	err = key.SetDWordValue("ConnectToRdp", config.ConnectToRdp)
	if err != nil {
		return err
	}

	return nil
}

// EnableSoftwareSASGeneration 启用软件SAS模拟输入
func EnableSoftwareSASGeneration() error {
	// SAS模拟输入
	key, err := registry.OpenKey(registry.LOCAL_MACHINE, "SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Policies\\System", registry.ALL_ACCESS)
	if err != nil {
		if err == registry.ErrNotExist {
			// 项不存在 创建它
			key, _, err = registry.CreateKey(registry.LOCAL_MACHINE, "SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Policies\\System", registry.ALL_ACCESS)
			if err != nil {
				return err
			}
		} else {
			return err
		}
	}

	defer key.Close()

	err = key.SetDWordValue("SoftwareSASGeneration", 1)
	if err != nil {
		return err
	}

	return nil
}

// EnableWinlogonDisableCAD 关闭 Ctrl + Alt + Del 组合键来登录
func EnableWinlogonDisableCAD() error {
	key, err := registry.OpenKey(registry.LOCAL_MACHINE, "SOFTWARE\\Microsoft\\Windows NT\\CurrentVersion\\Winlogon", registry.ALL_ACCESS)
	if err != nil {
		if err == registry.ErrNotExist {
			// 项不存在 创建它
			key, _, err = registry.CreateKey(registry.LOCAL_MACHINE, "SOFTWARE\\Microsoft\\Windows NT\\CurrentVersion\\Winlogon", registry.ALL_ACCESS)
			if err != nil {
				return err
			}
		} else {
			return err
		}
	}

	defer key.Close()

	err = key.SetDWordValue("DisableCAD", 1)
	if err != nil {
		return err
	}

	return nil
}

// EnableSystemDisableCAD 关闭 Ctrl + Alt + Del 组合键来登录
func EnableSystemDisableCAD() error {
	key, err := registry.OpenKey(registry.LOCAL_MACHINE, "SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Policies\\System", registry.ALL_ACCESS)
	if err != nil {
		if err == registry.ErrNotExist {
			// 项不存在 创建它
			key, _, err = registry.CreateKey(registry.LOCAL_MACHINE, "SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\Policies\\System", registry.ALL_ACCESS)
			if err != nil {
				return err
			}
		} else {
			return err
		}
	}

	defer key.Close()

	err = key.SetDWordValue("DisableCAD", 1)
	if err != nil {
		return err
	}

	return nil
}

// createRandomString 生成指定长度的随机字符串
func createRandomString(len int) string {
	var container string
	var str = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890"
	b := bytes.NewBufferString(str)
	length := b.Len()
	bigInt := big.NewInt(int64(length))
	for i := 0; i < len; i++ {
		randomInt, _ := rand.Int(rand.Reader, bigInt)
		container += string(str[randomInt.Int64()])
	}
	return container
}