fleetforge/agents/forge-agent/internal/vnc/vnc.go

package vnc

import (
	"crypto/des"
	"encoding/hex"
	"fmt"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"time"

	"github.com/fleetforge/forge-agent/internal/pkgmanager"
)

const waitForPortTimeout = 25 * time.Second

type Server struct {
	port     int
	pid      int
	passFile string
}

// Ensure starts a VNC server on the given port with the given password and
// returns a handle that can be used to stop it again.
func Ensure(port int, password string) (*Server, error) {
	switch runtime.GOOS {
	case "windows":
		return startTightVNC(port, password)
	case "linux":
		return startX11VNC(port, password)
	default:
		return nil, fmt.Errorf("remote desktop is not supported on %s yet", runtime.GOOS)
	}
}

func (s *Server) Stop() {
	if s == nil {
		return
	}
	switch runtime.GOOS {
	case "windows":
		stopTightVNC(s)
	default:
		stopX11VNC(s)
	}
}

func (s *Server) Port() int {
	return s.port
}

// ---------------------------------------------------------------- Windows

func tightVNCPath() string {
	for _, base := range []string{
		os.ExpandEnv(`%ProgramFiles%\TightVNC\tvnserver.exe`),
		os.ExpandEnv(`%ProgramFiles(x86)%\TightVNC\tvnserver.exe`),
		`C:\Program Files\TightVNC\tvnserver.exe`,
		`C:\Program Files (x86)\TightVNC\tvnserver.exe`,
	} {
		if _, err := os.Stat(base); err == nil {
			return base
		}
	}
	return ""
}

func quickRegAdd(subkey, name, value, regType string) error {
	cmd := exec.Command("reg", "add", subkey, "/v", name, "/t", regType, "/d", value, "/f")
	if out, err := cmd.CombinedOutput(); err != nil {
		return fmt.Errorf("reg add %s: %w (%s)", name, err, strings.TrimSpace(string(out)))
	}
	return nil
}

// configTightVNC writes the TightVNC server settings to both the global (HKLM)
// and the current user (HKCU) registry locations. Application mode reads the
// user settings, while the (optional) service mode reads the global ones, so we
// configure both. HKCU is written as the interactive user because reg.exe under
// the SYSTEM account would target the system profile instead.
func configTightVNC(port int, enc []byte) error {
	type entry struct {
		name string
		val  string
		typ  string
	}
	entries := []entry{
		{"RfbPort", strconv.Itoa(port), "REG_DWORD"},
		{"UsePasswordAuthentication", "1", "REG_DWORD"},
		{"AcceptLoopbackConnections", "1", "REG_DWORD"},
		{"Password", hex.EncodeToString(enc), "REG_BINARY"},
	}

	const globalKey = `HKLM\Software\TightVNC\Server`
	const userKey = `HKCU\Software\TightVNC\Server`

	write := func(subkey string, inUserSession bool) error {
		for _, e := range entries {
			args := []string{"add", subkey, "/v", e.name, "/t", e.typ, "/d", e.val, "/f"}
			var err error
			if inUserSession {
				_, code, uerr := pkgmanager.RunInUserSessionCapture("reg", args)
				if uerr != nil {
					err = uerr
				} else if code != 0 {
					err = fmt.Errorf("reg add %s (as user) exited %d", e.name, code)
				}
			} else {
				err = quickRegAdd(subkey, e.name, e.val, e.typ)
			}
			if err != nil {
				return fmt.Errorf("set %s\\%s: %w", subkey, e.name, err)
			}
		}
		return nil
	}

	if err := write(globalKey, false); err != nil {
		return err
	}
	if err := write(userKey, true); err != nil {
		return err
	}
	return nil
}

func startTightVNC(port int, password string) (*Server, error) {
	exe := tightVNCPath()
	if exe == "" {
		if _, err := pkgmanager.InstallPackageWithManager("TightVNC.TightVNC", "winget", ""); err != nil {
			if _, ferr := pkgmanager.InstallPackageWithManager("tightvnc", "choco", ""); ferr != nil {
				return nil, fmt.Errorf("tightvnc not installed and automatic install failed: %v (fallback %v)", err, ferr)
			}
		}
		exe = tightVNCPath()
		if exe == "" {
			return nil, fmt.Errorf("failed to locate tvnserver.exe after installation")
		}
	}

	// Best effort: stop any TightVNC service so it does not fight on the port.
	stopTightVNCServices()

	enc, err := encryptPassword(password)
	if err != nil {
		return nil, fmt.Errorf("vnc password encryption failed: %w", err)
	}

	if err := configTightVNC(port, enc); err != nil {
		return nil, err
	}

	// Start in the interactive user's session so the real desktop is served.
	pid, err := pkgmanager.LaunchInUserSession(exe, []string{"-run"})
	if err != nil {
		return nil, fmt.Errorf("failed to start tvnserver in user session: %w", err)
	}

	srv := &Server{port: port, pid: int(pid)}
	if err := waitForPort(port); err != nil {
		srv.Stop()
		return nil, err
	}
	return srv, nil
}

func stopTightVNCServices() {
	for _, name := range []string{"tvnserver", "TightVNC"} {
		_ = exec.Command("sc", "stop", name).Run()
	}
}

func stopTightVNC(s *Server) {
	stopTightVNCServices()
	if s.pid > 0 {
		if proc, err := os.FindProcess(s.pid); err == nil {
			_ = proc.Kill()
		}
	}
}

// encryptPassword produces the DES-encrypted 8-byte form used by the VNC
// servers to store the password (fixed key with bit-reversed bytes).
func encryptPassword(password string) ([]byte, error) {
	buf := make([]byte, 8)
	copy(buf, []byte(password))

	key := []byte{0xe8, 0x4a, 0xd6, 0x60, 0xc4, 0x72, 0x1a, 0xe3}
	for i := range key {
		key[i] = bitReverse(key[i])
	}

	block, err := des.NewCipher(key)
	if err != nil {
		return nil, err
	}

	out := make([]byte, 8)
	block.Encrypt(out, buf)
	return out, nil
}

func bitReverse(b byte) byte {
	var r byte
	for i := 0; i < 8; i++ {
		r = r<<1 | (b & 1)
		b >>= 1
	}
	return r
}

// ---------------------------------------------------------------- Linux

func startX11VNC(port int, password string) (*Server, error) {
	exe, err := exec.LookPath("x11vnc")
	if err != nil {
		if _, ierr := pkgmanager.InstallPackage("x11vnc"); ierr == nil {
			exe, err = exec.LookPath("x11vnc")
		}
		if err != nil {
			return nil, fmt.Errorf("x11vnc is not installed and could not be installed")
		}
	}

	display := os.Getenv("DISPLAY")
	if display == "" {
		display = ":0"
	}

	passFile := filepath.Join(os.TempDir(), fmt.Sprintf("ff-vnc-pass-%d.txt", time.Now().UnixNano()))
	if out, err := exec.Command(exe, "-storepasswd", password, passFile).CombinedOutput(); err != nil {
		return nil, fmt.Errorf("x11vnc -storepasswd failed: %w (%s)", err, strings.TrimSpace(string(out)))
	}

	cmd := exec.Command(
		exe,
		"-display", display,
		"-rfbport", strconv.Itoa(port),
		"-rfbauth", passFile,
		"-forever",
		"-shared",
	)
	if err := cmd.Start(); err != nil {
		os.Remove(passFile)
		return nil, fmt.Errorf("failed to start x11vnc: %w", err)
	}

	srv := &Server{port: port, pid: cmd.Process.Pid}
	srv.passFile = passFile
	if err := waitForPort(port); err != nil {
		srv.Stop()
		return nil, err
	}
	return srv, nil
}

func stopX11VNC(s *Server) {
	if s.pid > 0 {
		if proc, err := os.FindProcess(s.pid); err == nil {
			_ = proc.Kill()
		}
	}
	if s.passFile != "" {
		os.Remove(s.passFile)
	}
}

func waitForPort(port int) error {
	deadline := time.Now().Add(waitForPortTimeout)
	for time.Now().Before(deadline) {
		conn, err := net.DialTimeout("tcp", fmt.Sprintf("127.0.0.1:%d", port), 2*time.Second)
		if err == nil {
			conn.Close()
			return nil
		}
		time.Sleep(1 * time.Second)
	}
	return fmt.Errorf("VNC server did not start listening on port %d", port)
}