feat(server): add screenshot command

go.mod

@@ -6,6 +6,7 @@ toolchain go1.26.4
 
 require (
 	github.com/BurntSushi/toml v1.6.0
+	github.com/alexsnet/go-vnc v0.1.1-0.20230622101630-c02fb4ae6247 // We need an unreleased fix
 	github.com/cheggaaa/pb/v3 v3.1.7
 	github.com/dustin/go-humanize v1.0.1
 	github.com/fatih/color v1.19.0

go.sum

@@ -2,6 +2,8 @@ github.com/BurntSushi/toml v1.6.0 h1:dRaEfpa2VI55EwlIW72hMRHdWouJeRF7TPYhI+AUQjk
 github.com/BurntSushi/toml v1.6.0/go.mod h1:ukJfTF/6rtPPRCnwkur4qwRxa8vTRFBF0uk2lLoLwho=
 github.com/VividCortex/ewma v1.2.0 h1:f58SaIzcDXrSy3kWaHNvuJgJ3Nmz59Zji6XoJR/q1ow=
 github.com/VividCortex/ewma v1.2.0/go.mod h1:nz4BbCtbLyFDeC9SUHbtcT5644juEuWfUAUnGx7j5l4=
+github.com/alexsnet/go-vnc v0.1.1-0.20230622101630-c02fb4ae6247 h1:3Nb6ShfCJKO4WQ42ZyLL7ZZCjUjyi1/kP4CMQkmv1Kc=
+github.com/alexsnet/go-vnc v0.1.1-0.20230622101630-c02fb4ae6247/go.mod h1:bbRsg41Sh3zvrnWsw+REKJVGZd8Of2+S0V1G0ZaBhlU=
 github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
 github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw=
 github.com/bool64/dev v0.2.43 h1:yQ7qiZVef6WtCl2vDYU0Y+qSq+0aBrQzY8KXkklk9cQ=

internal/cmd/server/screenshot.go

@@ -0,0 +1,54 @@
+package server
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/spf13/cobra"
+
+	"github.com/hetznercloud/cli/internal/cmd/base"
+	"github.com/hetznercloud/cli/internal/cmd/cmpl"
+	"github.com/hetznercloud/cli/internal/cmd/server/screenshot"
+	"github.com/hetznercloud/cli/internal/hcapi2"
+	"github.com/hetznercloud/cli/internal/state"
+)
+
+var ScreenshotCmd = base.Cmd{
+	BaseCobraCommand: func(client hcapi2.Client) *cobra.Command {
+		cmd := &cobra.Command{
+			Use:                   "screenshot [options] <server>",
+			Short:                 "Take a screenshot of the Server VNC console",
+			ValidArgsFunction:     cmpl.SuggestArgs(cmpl.SuggestCandidatesF(client.Server().Names)),
+			TraverseChildren:      true,
+			DisableFlagsInUseLine: true,
+		}
+		cmd.Flags().StringP("output", "o", "screenshot.png", "Output file for the screenshot (must be a .png file).")
+		return cmd
+	},
+	Run: func(s state.State, cmd *cobra.Command, args []string) error {
+		idOrName := args[0]
+		server, _, err := s.Client().Server().Get(s, idOrName)
+		if err != nil {
+			return err
+		}
+		if server == nil {
+			return fmt.Errorf("Server not found: %s", idOrName)
+		}
+
+		result, _, err := s.Client().Server().RequestConsole(s, server)
+		if err != nil {
+			return err
+		}
+
+		if err := s.WaitForActions(s, cmd, result.Action); err != nil {
+			return err
+		}
+
+		output, _ := cmd.Flags().GetString("output")
+		if !strings.HasSuffix(strings.ToLower(output), ".png") {
+			cmd.Printf("Warning: Only .png output file are supported\n")
+		}
+
+		return screenshot.TakeScreenshot(s, result.WSSURL, output)
+	},
+}

internal/cmd/server/screenshot/screenshot.go

@@ -0,0 +1,139 @@
+package screenshot
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"image"
+	"image/color"
+	"image/png"
+	"io"
+	"log/slog"
+	"net/url"
+	"os"
+
+	"github.com/alexsnet/go-vnc"
+	"github.com/alexsnet/go-vnc/rfbflags"
+	"golang.org/x/net/websocket"
+)
+
+func websocketOrigin(wsURL string) (string, error) {
+	u, err := url.Parse(wsURL)
+	if err != nil {
+		return "", err
+	}
+	u.Scheme = "https"
+	return u.String(), nil
+}
+
+func dialWebsocket(ctx context.Context, wsURL string) (*websocket.Conn, error) {
+	origin, err := websocketOrigin(wsURL)
+	if err != nil {
+		return nil, err
+	}
+
+	cfg, err := websocket.NewConfig(wsURL, origin)
+	if err != nil {
+		return nil, err
+	}
+
+	ws, err := cfg.DialContext(ctx)
+	if err != nil {
+		return nil, err
+	}
+	ws.PayloadType = websocket.BinaryFrame
+
+	return ws, nil
+}
+
+func TakeScreenshot(ctx context.Context, wsURL string, filename string) error {
+	slog.Debug("dialing websocket")
+	ws, err := dialWebsocket(ctx, wsURL)
+	if err != nil {
+		return err
+	}
+	defer ws.Close()
+
+	slog.Debug("creating vnc client")
+	vncCfg := &vnc.ClientConfig{
+		Auth: []vnc.ClientAuth{
+			// Auth is managed at the HTTP/websocket level,therefore no password is
+			// needed for the VNC client.
+			&vnc.ClientAuthNone{},
+		},
+		ServerMessages: []vnc.ServerMessage{
+			&vnc.FramebufferUpdate{},
+		},
+		ServerMessageCh: make(chan vnc.ServerMessage, 1),
+	}
+
+	slog.Debug("connecting vnc client")
+	vncConn, err := vnc.Connect(ctx, ws, vncCfg)
+	if err != nil {
+		return err
+	}
+	defer vncConn.Close()
+
+	go func() {
+		slog.Debug("listening for vnc server messages")
+		if err := vncConn.ListenAndHandle(); err != nil {
+			slog.Error("failed to listen and handle incoming vnc server messages", "err", err)
+		}
+	}()
+
+	slog.Debug("requesting framebuffer update")
+	if err := vncConn.FramebufferUpdateRequest(
+		rfbflags.RFBFalse,
+		0, 0,
+		vncConn.FramebufferWidth(), vncConn.FramebufferHeight(),
+	); err != nil {
+		return err
+	}
+
+	slog.Debug("waiting for server message")
+	msg := <-vncCfg.ServerMessageCh
+
+	framebufferUpdate, ok := msg.(*vnc.FramebufferUpdate)
+	if !ok {
+		return fmt.Errorf("received unexpected server message: %T", msg)
+	}
+
+	// Sanity checks
+	if len(framebufferUpdate.Rects) != 1 ||
+		framebufferUpdate.Rects[0].X != 0 ||
+		framebufferUpdate.Rects[0].Y != 0 ||
+		framebufferUpdate.Rects[0].Width != vncConn.FramebufferWidth() ||
+		framebufferUpdate.Rects[0].Height != vncConn.FramebufferHeight() {
+		return fmt.Errorf("received invalid frame buffer update")
+	}
+
+	width := int(vncConn.FramebufferWidth())
+	height := int(vncConn.FramebufferHeight())
+
+	slog.Debug("composing image from framebuffer")
+	img := image.NewRGBA(image.Rect(0, 0, width, height))
+	{
+		rect := framebufferUpdate.Rects[0]
+		enc := rect.Enc.(*vnc.RawEncoding)
+		for index, clr := range enc.Colors {
+			x, y := index%width, index/width
+			img.Set(int(rect.X)+x, int(rect.Y)+y, color.RGBA{uint8(clr.R), uint8(clr.G), uint8(clr.B), 255})
+		}
+	}
+
+	slog.Debug("encoding image")
+	imgData := bytes.NewBuffer(nil)
+	{
+		pngEncoder := png.Encoder{CompressionLevel: png.DefaultCompression}
+		if err := pngEncoder.Encode(io.Writer(imgData), img); err != nil {
+			return err
+		}
+	}
+
+	slog.Debug("writing image to file")
+	if err := os.WriteFile(filename, imgData.Bytes(), 0600); err != nil {
+		return err
+	}
+
+	return nil
+}

internal/cmd/server/screenshot/screenshot_test.go

@@ -0,0 +1,37 @@
+package screenshot
+
+import (
+	"context"
+	"log/slog"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/hetznercloud/hcloud-go/v2/hcloud"
+)
+
+func TestTakeScreenshot(t *testing.T) {
+	t.Skip()
+
+	ctx := t.Context()
+
+	// TODO: Load token
+	client := hcloud.NewClient(hcloud.WithToken("TOKEN"))
+
+	// TODO: Create or use existing server
+	result, _, err := client.Server.RequestConsole(ctx, &hcloud.Server{ID: 0})
+	require.NoError(t, err)
+
+	err = client.Action.WaitFor(ctx, result.Action)
+	require.NoError(t, err)
+
+	ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second)
+	defer cancel()
+
+	slog.SetDefault(slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelDebug})))
+
+	err = TakeScreenshot(ctx, result.WSSURL, "screenshot.png")
+	require.NoError(t, err)
+}

internal/cmd/server/server.go

@@ -73,6 +73,7 @@ func NewCommand(s state.State) *cobra.Command {
 		SSHCmd.CobraCommand(s),
 		IPCmd.CobraCommand(s),
 		RequestConsoleCmd.CobraCommand(s),
+		ScreenshotCmd.CobraCommand(s),
 		ResetPasswordCmd.CobraCommand(s),
 		MetricsCmd.CobraCommand(s),
 		SetRDNSCmd.CobraCommand(s),