From 885dcb8ceeb3c0be415b2203dcd8c4d34b5476c2 Mon Sep 17 00:00:00 2001 From: jo Date: Mon, 15 Jun 2026 15:20:42 +0200 Subject: [PATCH] feat(server): add screenshot command Adds a new `hcloud server screenshot -o screenshot.png `, that take a screenshot of the server console using a VNC client. --- docs/reference/manual/hcloud_server.md | 1 + .../manual/hcloud_server_screenshot.md | 33 +++++ go.mod | 1 + go.sum | 2 + internal/cmd/server/screenshot.go | 54 +++++++ internal/cmd/server/screenshot/screenshot.go | 139 ++++++++++++++++++ .../cmd/server/screenshot/screenshot_test.go | 37 +++++ internal/cmd/server/server.go | 1 + 8 files changed, 268 insertions(+) create mode 100644 docs/reference/manual/hcloud_server_screenshot.md create mode 100644 internal/cmd/server/screenshot.go create mode 100644 internal/cmd/server/screenshot/screenshot.go create mode 100644 internal/cmd/server/screenshot/screenshot_test.go diff --git a/docs/reference/manual/hcloud_server.md b/docs/reference/manual/hcloud_server.md index 054b0a04c..f111b2831 100644 --- a/docs/reference/manual/hcloud_server.md +++ b/docs/reference/manual/hcloud_server.md @@ -55,6 +55,7 @@ Manage Servers * [hcloud server request-console](hcloud_server_request-console.md) - Request a WebSocket VNC console for a Server * [hcloud server reset](hcloud_server_reset.md) - Reset a Server * [hcloud server reset-password](hcloud_server_reset-password.md) - Reset the root password of a Server +* [hcloud server screenshot](hcloud_server_screenshot.md) - Take a screenshot of the Server VNC console * [hcloud server set-rdns](hcloud_server_set-rdns.md) - Change reverse DNS of a Server * [hcloud server shutdown](hcloud_server_shutdown.md) - Shutdown a server * [hcloud server ssh](hcloud_server_ssh.md) - Spawn an SSH connection for the Server diff --git a/docs/reference/manual/hcloud_server_screenshot.md b/docs/reference/manual/hcloud_server_screenshot.md new file mode 100644 index 000000000..689154adc --- /dev/null +++ b/docs/reference/manual/hcloud_server_screenshot.md @@ -0,0 +1,33 @@ +## hcloud server screenshot + +Take a screenshot of the Server VNC console + +``` +hcloud server screenshot [options] +``` + +### Options + +``` + -h, --help help for screenshot + -o, --output string Output file for the screenshot (must be a .png file). (default "screenshot.png") +``` + +### Options inherited from parent commands + +``` + --config string Config file path (default "~/.config/hcloud/cli.toml") + --context string Currently active context + --debug Enable debug output + --debug-file string File to write debug output to + --endpoint string Hetzner Cloud API endpoint (default "https://api.hetzner.cloud/v1") + --hetzner-endpoint string Hetzner API endpoint (default "https://api.hetzner.com/v1") + --no-experimental-warnings If true, experimental warnings are not shown + --poll-interval duration Interval at which to poll information, for example action progress (default 500ms) + --quiet If true, only print error messages +``` + +### SEE ALSO + +* [hcloud server](hcloud_server.md) - Manage Servers + diff --git a/go.mod b/go.mod index f8e0add91..505b184f7 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index a53d3680d..18cdf94a2 100644 --- a/go.sum +++ b/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= diff --git a/internal/cmd/server/screenshot.go b/internal/cmd/server/screenshot.go new file mode 100644 index 000000000..a54c95a0b --- /dev/null +++ b/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] ", + 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) + }, +} diff --git a/internal/cmd/server/screenshot/screenshot.go b/internal/cmd/server/screenshot/screenshot.go new file mode 100644 index 000000000..29cc23db5 --- /dev/null +++ b/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 +} diff --git a/internal/cmd/server/screenshot/screenshot_test.go b/internal/cmd/server/screenshot/screenshot_test.go new file mode 100644 index 000000000..b2936b680 --- /dev/null +++ b/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) +} diff --git a/internal/cmd/server/server.go b/internal/cmd/server/server.go index 7fdf68eb8..29a23a48c 100644 --- a/internal/cmd/server/server.go +++ b/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),