diff --git a/README.md b/README.md index bab03d4..77ca38a 100644 --- a/README.md +++ b/README.md @@ -91,6 +91,7 @@ DUT agent on it. Below you can find the currently supported modules with example | [GPIO Button](./pkg/module/gpio/README.md) | :white_check_mark: | | [GPIO Switch](./pkg/module/gpio/README.md) | :white_check_mark: | | [IPMI Power Control](./pkg/module/ipmi/README.md) | :white_check_mark: | +| [PiKVM](./pkg/module/pikvm/README.md) | :white_check_mark: | | [Power Distribution Unit (Intellinet)](./pkg/module/pdu/README.md) | :white_check_mark: | | Power Distribution Unit (Delock) | :hourglass_flowing_sand: | | [SPI Flash Emulator](./pkg/module/flash-emulate/README.md) | :white_check_mark: | diff --git a/cmds/dutagent/modules.go b/cmds/dutagent/modules.go index 61999df..e47ca1c 100644 --- a/cmds/dutagent/modules.go +++ b/cmds/dutagent/modules.go @@ -20,6 +20,7 @@ import ( _ "github.com/BlindspotSoftware/dutctl/pkg/module/gpio" _ "github.com/BlindspotSoftware/dutctl/pkg/module/ipmi" _ "github.com/BlindspotSoftware/dutctl/pkg/module/pdu" + _ "github.com/BlindspotSoftware/dutctl/pkg/module/pikvm" _ "github.com/BlindspotSoftware/dutctl/pkg/module/serial" _ "github.com/BlindspotSoftware/dutctl/pkg/module/shell" _ "github.com/BlindspotSoftware/dutctl/pkg/module/ssh" diff --git a/cmds/dutctl/dutctl.go b/cmds/dutctl/dutctl.go index 3810743..7d786bd 100644 --- a/cmds/dutctl/dutctl.go +++ b/cmds/dutctl/dutctl.go @@ -7,6 +7,7 @@ package main import ( + "crypto/sha256" "crypto/tls" "errors" "flag" @@ -16,6 +17,8 @@ import ( "net" "net/http" "os" + "strconv" + "strings" "connectrpc.com/connect" "github.com/BlindspotSoftware/dutctl/internal/buildinfo" @@ -176,16 +179,18 @@ func (app *application) start() { app.exit(errInvalidCmdline) } - err := app.listRPC() - app.exit(err) + app.exit(app.listRPC()) } if len(app.args) == 1 { - device := app.args[0] - err := app.commandsRPC(device) - app.exit(err) + app.exit(app.commandsRPC(app.args[0])) } + app.dispatchDeviceCommand() +} + +// dispatchDeviceCommand handles the "device command [args...]" invocation form. +func (app *application) dispatchDeviceCommand() { device := app.args[0] command := app.args[1] cmdArgs := app.args[2:] @@ -198,12 +203,16 @@ func (app *application) start() { } if len(cmdArgs) > 0 && cmdArgs[0] == "help" { - err := app.detailsRPC(device, command, "help") + app.exit(app.detailsRPC(device, command, "help")) + } + + // Preprocess arguments for optimization (e.g., calculate hashes) + cmdArgs, err := app.preprocessArgs(command, cmdArgs) + if err != nil { app.exit(err) } - err := app.runRPC(device, command, cmdArgs) - app.exit(err) + app.exit(app.runRPC(device, command, cmdArgs)) } // exit terminates the application. If the provided error is not nil, it is printed to @@ -235,6 +244,59 @@ func (app *application) exit(err error) { app.exitFunc(1) } +// preprocessArgs preprocesses command arguments for optimization. +// For example, it calculates file hashes for PiKVM media mount commands. +func (app *application) preprocessArgs(command string, args []string) ([]string, error) { + // Argument counts for "media mount [hash] [size]". + const ( + mediaMountMinArgs = 2 // mount + path + mediaMountWithHash = 3 // mount + path + hash + ) + + // Optimize PiKVM media mount: calculate hash locally to avoid unnecessary transfers + if strings.ToLower(command) == "media" && len(args) >= mediaMountMinArgs && strings.ToLower(args[0]) == "mount" { + imagePath := args[1] + + // If hash is already provided (args[2]), skip preprocessing + if len(args) >= mediaMountWithHash { + return args, nil + } + + // Check if file exists + fileInfo, err := os.Stat(imagePath) + if err != nil { + // If file doesn't exist locally, let the agent handle the error + return args, nil + } + + // Calculate SHA256 hash + log.Printf("Calculating SHA256 hash of %s...", imagePath) + + file, err := os.Open(imagePath) + if err != nil { + return args, fmt.Errorf("failed to open file for hashing: %w", err) + } + defer file.Close() + + hash := sha256.New() + + _, err = io.Copy(hash, file) + if err != nil { + return args, fmt.Errorf("failed to calculate hash: %w", err) + } + + hashSum := fmt.Sprintf("%x", hash.Sum(nil)) + fileSize := fileInfo.Size() + + log.Printf("Hash: %s, Size: %d bytes", hashSum, fileSize) + + // Append hash and size to arguments + return append(args, hashSum, strconv.FormatInt(fileSize, 10)), nil + } + + return args, nil +} + func (app *application) printVersion() { app.formatter.WriteContent(output.Content{ Type: output.TypeVersion, diff --git a/pkg/module/pikvm/README.md b/pkg/module/pikvm/README.md new file mode 100644 index 0000000..458feb9 --- /dev/null +++ b/pkg/module/pikvm/README.md @@ -0,0 +1,97 @@ +# PiKVM + +This module provides comprehensive control of a DUT via a PiKVM device. It offers power management through ATX control, keyboard input simulation, and virtual media mounting capabilities. + +PiKVM documentation: https://docs.pikvm.org/ + +Compatibility: This module uses PiKVM's HTTP API endpoints (e.g. `/api/atx/*`, `/api/hid/*`, `/api/msd/*`, `/api/streamer/snapshot`). It is intended for PiKVM (kvmd/PiKVM OS) systems, and should work with other devices only if they expose a compatible API. + +## Features + +### Power Management +Control the DUT's power state via ATX power and reset buttons: + +- Power commands use PiKVM's ATX control API: + - `on` is idempotent (does nothing if already powered on) + - `off` performs a graceful shutdown via power button press + - `force-off` performs a hard power off via long press (5+ seconds) + - `reset` triggers an ATX reset button press + - `force-reset` triggers a hardware hot reset + +``` +COMMANDS: + on Power on (does nothing if already on) + off Graceful shutdown (soft power-off) + force-off Force power off (hard shutdown, 5+ second press) + reset Reset via ATX reset button + force-reset Force reset (hardware hot reset) + status Query current power state +``` + +### Keyboard Control +Send keyboard input to the DUT: + +See [pikvm-key-strings.md](./pikvm-key-strings.md) for the full list of supported key strings. + +``` +COMMANDS: + type Type a text string + key Send a single key (e.g., Enter, Escape, F12) + key-combo Send key combination (e.g., Ctrl+Alt+Delete) + +FLAGS (must come before the action): + --delay Delay between key events for key-combo (default: 500ms) +``` + +### Virtual Media +Mount ISO images or disk images as virtual USB devices: + +- `mount` uploads images to PiKVM's storage (with automatic space management) +- `mount-url` instructs PiKVM to download the image directly from the URL +- Old images are automatically deleted when storage space is needed +- SHA256 hashing prevents duplicate uploads + +``` +COMMANDS: + mount Mount an image file from the agent's filesystem + mount-url Mount an image from a URL + unmount Unmount current virtual media + media-status Show mounted media information +``` + +### Screenshot +Capture the current display output from PiKVM's video stream. + +## Configuration Options + +| Option | Type | Default | Description | +| -------- | ------ | ------- | ---------------------------------------------------------------------------------------------------------------------------------------------------- | +| host | string | - | Address of the PiKVM device (e.g., "192.168.1.100") Defaults to HTTPS if no scheme is specified, HTTP can be used by explicitly specifying `http://` | +| user | string | admin | Username for authentication | +| password | string | - | Password for authentication | +| timeout | string | 10s | Timeout for HTTP requests (e.g., "10s", "30s") | +| mode | string | - | **Required**: Mode ("power", "keyboard", "media", "screenshot") | + +⚠️ **Security Warning**: Passwords are stored in plaintext in the configuration file. This should only be used in trusted environments. + +## Usage Examples + +See [pikvm-example-cfg.yml](./pikvm-example-cfg.yml) for comprehensive configuration examples. + +### Basic Power Control +See [pikvm-example-power.yml](./pikvm-example-power.yml). + +### Boot Menu Access +See [pikvm-example-keyboard.yml](./pikvm-example-keyboard.yml). + +### ISO Mounting +See [pikvm-example-media.yml](./pikvm-example-media.yml). + +### Screenshot Capture +See [pikvm-example-screenshot.yml](./pikvm-example-screenshot.yml). + +## Requirements + +- PiKVM device with API access enabled +- Network connectivity between dutagent and PiKVM +- Valid authentication credentials diff --git a/pkg/module/pikvm/keyboard.go b/pkg/module/pikvm/keyboard.go new file mode 100644 index 0000000..b2b521b --- /dev/null +++ b/pkg/module/pikvm/keyboard.go @@ -0,0 +1,386 @@ +// Copyright 2025 Blindspot Software +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package pikvm + +import ( + "bytes" + "context" + "fmt" + "net/http" + "net/url" + "strings" + "time" + + "github.com/BlindspotSoftware/dutctl/pkg/module" +) + +// Keyboard command constants. +const ( + typeCmd = "type" + key = "key" + keyCombo = "key-combo" +) + +const ( + maxTextLength = 1024 + minComboKeys = 2 + delayFlagArgCount = 2 // "--delay" flag plus its value + invalidCharNul = "\x00" + invalidCharNewline = "\n" + invalidCharReturn = "\r" + defaultKeyDelay = 500 * time.Millisecond +) + +func parseKeyboardFlags(args []string) (time.Duration, []string, error) { + delay := defaultKeyDelay + + for len(args) > 0 { + arg := args[0] + if !strings.HasPrefix(arg, "--") { + break + } + + if strings.HasPrefix(arg, "--delay=") { + value := strings.TrimPrefix(arg, "--delay=") + + parsed, parseErr := time.ParseDuration(value) + if parseErr != nil || parsed <= 0 { + return 0, nil, fmt.Errorf("invalid --delay value %q", value) + } + + delay = parsed + args = args[1:] + + continue + } + + if arg == "--delay" { + if len(args) < delayFlagArgCount { + return 0, nil, fmt.Errorf("--delay requires a duration value") + } + + value := args[1] + + parsed, parseErr := time.ParseDuration(value) + if parseErr != nil || parsed <= 0 { + return 0, nil, fmt.Errorf("invalid --delay value %q", value) + } + + delay = parsed + args = args[2:] + + continue + } + + return 0, nil, fmt.Errorf("unknown flag %q", arg) + } + + return delay, args, nil +} + +// handleKeyboardCommandRouter routes keyboard commands based on the first argument. +func (p *PiKVM) handleKeyboardCommandRouter(ctx context.Context, s module.Session, args []string) error { + delay, remaining, err := parseKeyboardFlags(args) + if err != nil { + s.Printf("Error: %v\n", err) + s.Println("Usage: keyboard [--delay ] ...") + + return nil + } + + if len(remaining) == 0 { + s.Println("Keyboard command requires an action: type|key|key-combo") + + return nil + } + + command := strings.ToLower(remaining[0]) + + return p.handleKeyboardCommand(ctx, s, command, remaining, delay) +} + +// handleKeyboardCommand dispatches keyboard-related commands. +func (p *PiKVM) handleKeyboardCommand(ctx context.Context, s module.Session, command string, args []string, delay time.Duration) error { + switch command { + case typeCmd: + if len(args) < minArgsRequired { + s.Println("Error: 'type' command requires text argument") + + return nil + } + + return p.handleType(ctx, s, strings.Join(args[1:], " ")) + case key: + if len(args) < minArgsRequired { + s.Println("Error: 'key' command requires key name argument") + + return nil + } + + return p.handleKey(ctx, s, args[1]) + case keyCombo: + if len(args) < minArgsRequired { + s.Println("Error: 'key-combo' command requires key combination argument") + + return nil + } + + return p.handleCombo(ctx, s, args[1], delay) + default: + return fmt.Errorf("unknown keyboard action: %s (must be: type, key, key-combo)", command) + } +} + +func (p *PiKVM) handleType(ctx context.Context, s module.Session, text string) error { + // Validate text input + if text == "" { + return fmt.Errorf("empty text input") + } + + // Check for invalid characters + if idx := strings.IndexAny(text, invalidCharNul+invalidCharNewline+invalidCharReturn); idx != -1 { + return fmt.Errorf("text contains invalid character: %q", rune(text[idx])) + } + + // Limit text length + if len(text) > maxTextLength { + return fmt.Errorf("text input too long (%d chars, max is %d)", len(text), maxTextLength) + } + + s.Printf("Typing text: %s\n", text) + + // Use the /api/hid/print endpoint for text input + resp, err := p.doRequest(ctx, http.MethodPost, "/api/hid/print?slow=true", bytes.NewBufferString(text), "text/plain", requestOptions{}) + if err != nil { + return fmt.Errorf("failed to type text: %v", err) + } + defer resp.Body.Close() + + s.Println("Text typed successfully") + + return nil +} + +func (p *PiKVM) handleKey(ctx context.Context, s module.Session, keyName string) error { + // Map the key name to PiKVM format + mappedKey, err := mapToKeyboardKey(keyName) + if err != nil { + return fmt.Errorf("error with key '%s': %w", keyName, err) + } + + s.Printf("Sending key: %s\n", keyName) + + // Send key press (press and release) + err = p.sendKeyRequest(ctx, fmt.Sprintf("/api/hid/events/send_key?key=%s", url.QueryEscape(mappedKey))) + if err != nil { + return err + } + + s.Println("Key sent successfully") + + return nil +} + +func (p *PiKVM) handleCombo(ctx context.Context, s module.Session, comboStr string, delay time.Duration) error { + // Parse combo like "Ctrl+Alt+Delete" into array of keys + keys := strings.Split(comboStr, "+") + for idx := range keys { + keys[idx] = strings.TrimSpace(keys[idx]) + } + + if len(keys) < minComboKeys { + return fmt.Errorf("key combination must have at least %d keys", minComboKeys) + } + + s.Printf("Sending key combination: %s\n", comboStr) + + // Validate and map keys + modifiers := keys[:len(keys)-1] + mainKey := keys[len(keys)-1] + + mainKeyMapped, mappedModifiers, err := p.validateAndMapComboKeys(modifiers, mainKey) + if err != nil { + return err + } + + // Execute the key combination + err = p.executeKeyCombo(ctx, mappedModifiers, mainKeyMapped, delay) + if err != nil { + return err + } + + s.Println("Key combination sent successfully") + + return nil +} + +// validateAndMapComboKeys validates and maps the keys in a combination. +func (p *PiKVM) validateAndMapComboKeys(modifiers []string, mainKey string) (string, []string, error) { + // Map main key + mainKeyMapped, err := mapToKeyboardKey(mainKey) + if err != nil { + return "", nil, fmt.Errorf("error with key '%s': %w", mainKey, err) + } + + // Check if the last key is not a modifier + if isModifierKey(mainKeyMapped) { + return "", nil, fmt.Errorf("last key in combination should not be a modifier: %s", mainKey) + } + + // Map and validate modifiers + mappedModifiers := make([]string, len(modifiers)) + for idx, mod := range modifiers { + var err error + + mappedModifiers[idx], err = mapToKeyboardKey(mod) + if err != nil { + return "", nil, fmt.Errorf("error with modifier '%s': %w", mod, err) + } + + if !isModifierKey(mappedModifiers[idx]) { + return "", nil, fmt.Errorf("invalid modifier key: %s", mod) + } + } + + return mainKeyMapped, mappedModifiers, nil +} + +// executeKeyCombo executes a key combination by pressing modifiers, main key, and releasing modifiers. +func (p *PiKVM) executeKeyCombo(ctx context.Context, modifiers []string, mainKey string, delay time.Duration) error { + // Press down all modifier keys + for _, modifier := range modifiers { + err := p.sendKeyRequest(ctx, fmt.Sprintf("/api/hid/events/send_key?key=%s&state=1", url.QueryEscape(modifier))) + if err != nil { + return err + } + + time.Sleep(delay) + } + + // Press and release the main key + err := p.sendKeyRequest(ctx, fmt.Sprintf("/api/hid/events/send_key?key=%s", url.QueryEscape(mainKey))) + if err != nil { + return err + } + + time.Sleep(delay) + + // Release all modifier keys in reverse order + for idx := len(modifiers) - 1; idx >= 0; idx-- { + err := p.sendKeyRequest(ctx, fmt.Sprintf("/api/hid/events/send_key?key=%s&state=0", url.QueryEscape(modifiers[idx]))) + if err != nil { + return err + } + + time.Sleep(delay) + } + + return nil +} + +// sendKeyRequest sends a key request to PiKVM. +func (p *PiKVM) sendKeyRequest(ctx context.Context, urlPath string) error { + resp, err := p.doRequest(ctx, http.MethodPost, urlPath, nil, "", requestOptions{}) + if err != nil { + return err + } + defer resp.Body.Close() + + return nil +} + +// Map of keyboard key names to PiKVM web_name values. +// +//nolint:gochecknoglobals // Keyboard mapping table is a legitimate use case for a package-level variable +var keyboardMap = map[string]string{ + // Letter keys + "a": "KeyA", "b": "KeyB", "c": "KeyC", "d": "KeyD", "e": "KeyE", + "f": "KeyF", "g": "KeyG", "h": "KeyH", "i": "KeyI", "j": "KeyJ", + "k": "KeyK", "l": "KeyL", "m": "KeyM", "n": "KeyN", "o": "KeyO", + "p": "KeyP", "q": "KeyQ", "r": "KeyR", "s": "KeyS", "t": "KeyT", + "u": "KeyU", "v": "KeyV", "w": "KeyW", "x": "KeyX", "y": "KeyY", + "z": "KeyZ", + + // Digits + "0": "Digit0", "1": "Digit1", "2": "Digit2", "3": "Digit3", "4": "Digit4", + "5": "Digit5", "6": "Digit6", "7": "Digit7", "8": "Digit8", "9": "Digit9", + + // Modifiers + "ctrl": "ControlLeft", "control": "ControlLeft", + "alt": "AltLeft", + "shift": "ShiftLeft", + "meta": "MetaLeft", "win": "MetaLeft", "windows": "MetaLeft", "cmd": "MetaLeft", "command": "MetaLeft", + + // Common keys + "space": "Space", + "enter": "Enter", "return": "Enter", + "tab": "Tab", + "esc": "Escape", "escape": "Escape", + "backspace": "Backspace", + "delete": "Delete", "del": "Delete", + "insert": "Insert", "ins": "Insert", + "home": "Home", + "end": "End", + "pageup": "PageUp", "pgup": "PageUp", + "pagedown": "PageDown", "pgdn": "PageDown", + + // Function keys + "f1": "F1", "f2": "F2", "f3": "F3", "f4": "F4", "f5": "F5", "f6": "F6", + "f7": "F7", "f8": "F8", "f9": "F9", "f10": "F10", "f11": "F11", "f12": "F12", + + // Arrow keys + "up": "ArrowUp", "arrowup": "ArrowUp", + "down": "ArrowDown", "arrowdown": "ArrowDown", + "left": "ArrowLeft", "arrowleft": "ArrowLeft", + "right": "ArrowRight", "arrowright": "ArrowRight", + + // Punctuation and symbols + "-": "Minus", "minus": "Minus", + "=": "Equal", "equals": "Equal", + "[": "BracketLeft", "leftbracket": "BracketLeft", + "]": "BracketRight", "rightbracket": "BracketRight", + "\\": "Backslash", "backslash": "Backslash", + ";": "Semicolon", "semicolon": "Semicolon", + "'": "Quote", "quote": "Quote", + "`": "Backquote", "backquote": "Backquote", + ",": "Comma", "comma": "Comma", + ".": "Period", "period": "Period", + "/": "Slash", "slash": "Slash", + + // Additional keys + "capslock": "CapsLock", + "printscreen": "PrintScreen", + "pause": "Pause", + "scrolllock": "ScrollLock", "scroll": "ScrollLock", + "numlock": "NumLock", "num": "NumLock", + "contextmenu": "ContextMenu", "menu": "ContextMenu", + + // Numpad keys + "numpad0": "Numpad0", "numpad1": "Numpad1", "numpad2": "Numpad2", + "numpad3": "Numpad3", "numpad4": "Numpad4", "numpad5": "Numpad5", + "numpad6": "Numpad6", "numpad7": "Numpad7", "numpad8": "Numpad8", + "numpad9": "Numpad9", "numpad/": "NumpadDivide", "numpad*": "NumpadMultiply", + "numpad-": "NumpadSubtract", "numpad+": "NumpadAdd", + "numpadenter": "NumpadEnter", "numpad.": "NumpadDecimal", +} + +// mapToKeyboardKey maps common key names to PiKVM keyboard key names. +func mapToKeyboardKey(key string) (string, error) { + lowerKey := strings.ToLower(key) + + if value, exists := keyboardMap[lowerKey]; exists { + return value, nil + } + + return "", fmt.Errorf("unsupported key: %s", key) +} + +// isModifierKey checks if a key is a modifier key. +func isModifierKey(key string) bool { + return strings.HasPrefix(key, "Control") || + strings.HasPrefix(key, "Alt") || + strings.HasPrefix(key, "Shift") || + strings.HasPrefix(key, "Meta") +} diff --git a/pkg/module/pikvm/keyboard_test.go b/pkg/module/pikvm/keyboard_test.go new file mode 100644 index 0000000..53d31fc --- /dev/null +++ b/pkg/module/pikvm/keyboard_test.go @@ -0,0 +1,70 @@ +// Copyright 2025 Blindspot Software +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package pikvm + +import ( + "reflect" + "testing" + "time" +) + +func TestParseKeyboardFlags(t *testing.T) { + t.Parallel() + + testCases := []struct { + name string + args []string + wantDelay time.Duration + wantRemaining []string + wantErr bool + }{ + { + name: "default", + args: []string{"key", "F12"}, + wantDelay: defaultKeyDelay, + wantRemaining: []string{"key", "F12"}, + }, + { + name: "delay equals", + args: []string{"--delay=25ms", "key-combo", "Ctrl+Alt+Delete"}, + wantDelay: 25 * time.Millisecond, + wantRemaining: []string{"key-combo", "Ctrl+Alt+Delete"}, + }, + { + name: "delay separate", + args: []string{"--delay", "30ms", "key", "Enter"}, + wantDelay: 30 * time.Millisecond, + wantRemaining: []string{"key", "Enter"}, + }, + { + name: "invalid delay", + args: []string{"--delay", "nope", "key", "Enter"}, + wantErr: true, + }, + { + name: "unknown flag", + args: []string{"--wat", "key", "Enter"}, + wantErr: true, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + delay, remaining, err := parseKeyboardFlags(tc.args) + if (err != nil) != tc.wantErr { + t.Fatalf("error=%v, wantErr=%v", err, tc.wantErr) + } + if tc.wantErr { + return + } + if delay != tc.wantDelay { + t.Fatalf("delay=%v, want %v", delay, tc.wantDelay) + } + if !reflect.DeepEqual(remaining, tc.wantRemaining) { + t.Fatalf("remaining=%#v, want %#v", remaining, tc.wantRemaining) + } + }) + } +} diff --git a/pkg/module/pikvm/media.go b/pkg/module/pikvm/media.go new file mode 100644 index 0000000..66182fd --- /dev/null +++ b/pkg/module/pikvm/media.go @@ -0,0 +1,825 @@ +// Copyright 2025 Blindspot Software +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package pikvm + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/json" + "errors" + "fmt" + "io" + "mime/multipart" + "net/http" + "net/url" + "os" + "path/filepath" + "strconv" + "strings" + "time" + + "github.com/BlindspotSoftware/dutctl/pkg/module" +) + +// Media command constants. +const ( + mount = "mount" + mountURL = "mount-url" + unmount = "unmount" + mediaStatus = "media-status" + + mediaNone = "None" +) + +const ( + safetyMargin = 10 * 1024 * 1024 // 10 MB safety margin + scanningDelay = 5 * time.Second // Delay after deleting to allow PiKVM storage to update + maxDeletionRetries = 10 // Maximum number of images to delete before giving up + bytesPerMB = 1024 * 1024 // Bytes per megabyte + imageWaitTimeout = 60 * time.Second // Timeout waiting for image to be ready + imageWaitInterval = 1 * time.Second // Interval between image status checks + remoteDownloadTimout = 3600 // Timeout in seconds for remote image downloads +) + +var ErrMissingImage = errors.New("image not found in storage") + +// Status represents the PiKVM MSD status response. +type Status struct { + Ok bool `json:"ok"` + Result StatusResult `json:"result"` +} + +// StatusResult contains the actual status data. +type StatusResult struct { + Busy bool `json:"busy"` + Enabled bool `json:"enabled"` + Online bool `json:"online"` + Drive Drive `json:"drive"` + Storage struct { + Images map[string]StorageImage `json:"images"` + Parts map[string]StoragePart `json:"parts"` + } `json:"storage"` +} + +// Drive contains information about the virtual drive. +type Drive struct { + Connected bool `json:"connected"` + Image Image `json:"image"` + Cdrom bool `json:"cdrom"` + Rw bool `json:"rw"` +} + +// Image contains information about the mounted image. +type Image struct { + Name string `json:"name"` + Size uint64 `json:"size"` + Complete bool `json:"complete"` + InStorage bool `json:"in_storage"` //nolint:tagliatelle // PiKVM API uses snake_case +} + +// StorageImage contains information about an image in storage. +type StorageImage struct { + Complete bool `json:"complete"` + ModTS float32 `json:"mod_ts"` //nolint:tagliatelle // PiKVM API uses snake_case + Removable bool `json:"removable"` + Size uint64 `json:"size"` +} + +// StoragePart contains information about a storage partition. +type StoragePart struct { + Free uint64 `json:"free"` + Size uint64 `json:"size"` + Writable bool `json:"writable"` +} + +// handleMediaCommandRouter routes media commands based on the first argument. +func (p *PiKVM) handleMediaCommandRouter(ctx context.Context, s module.Session, args []string) error { + if len(args) == 0 { + s.Println("Media command requires an action: mount|mount-url|unmount|media-status") + + return nil + } + + command := strings.ToLower(args[0]) + + return p.handleMediaCommand(ctx, s, command, args) +} + +// handleMediaCommand dispatches virtual media commands. +func (p *PiKVM) handleMediaCommand(ctx context.Context, s module.Session, command string, args []string) error { + switch command { + case mount: + if len(args) < minArgsRequired { + s.Println("Error: 'mount' command requires file path argument") + + return nil + } + + precomputedHash, precomputedSize := parseMountArgs(args) + + return p.handleMount(ctx, s, args[1], precomputedHash, precomputedSize) + case mountURL: + if len(args) < minArgsRequired { + s.Println("Error: 'mount-url' command requires URL argument") + + return nil + } + + return p.handleMountURL(ctx, s, args[1]) + case unmount: + return p.handleUnmount(ctx, s) + case mediaStatus: + return p.handleMediaStatus(ctx, s) + default: + return fmt.Errorf("unknown media action: %s (must be: mount, mount-url, unmount, media-status)", command) + } +} + +// parseMountArgs extracts the optional precomputed hash and size from the +// "mount [hash] [size]" arguments. Missing values are returned as zero. +func parseMountArgs(args []string) (string, int64) { + const ( + hashArgIndex = 2 + sizeArgIndex = 3 + ) + + var ( + hash string + size int64 + ) + + if len(args) > hashArgIndex { + hash = args[hashArgIndex] + } + + if len(args) > sizeArgIndex { + parsed, err := strconv.ParseInt(args[sizeArgIndex], 10, 64) + if err == nil { + size = parsed + } + } + + return hash, size +} + +func (p *PiKVM) handleMount(ctx context.Context, s module.Session, imagePath, precomputedHash string, precomputedSize int64) error { + s.Printf("Preparing to mount image: %s\n", filepath.Base(imagePath)) + + // Fast path: if the client precomputed the hash and the image already + // exists on the PiKVM, mount it without transferring anything. + if precomputedHash != "" { + mounted, err := p.tryMountExisting(ctx, s, imagePath, precomputedHash) + if err != nil { + return err + } + + if mounted { + return nil + } + } + + return p.transferAndMount(ctx, s, imagePath, precomputedHash, precomputedSize) +} + +// tryMountExisting attempts to mount an image already present on the PiKVM, +// identified by its hash. It reports whether the image was found and mounted. +func (p *PiKVM) tryMountExisting(ctx context.Context, s module.Session, imagePath, hashSum string) (bool, error) { + s.Printf("Using precomputed hash: %s\n", hashSum) + + err := p.checkImageExists(ctx, hashSum) + if errors.Is(err, ErrMissingImage) { + s.Println("Image not found on PiKVM, will transfer from client.") + + return false, nil + } + + if err != nil { + return false, fmt.Errorf("failed to check if image exists: %v", err) + } + + s.Println("Image already exists on PiKVM, skipping upload.") + + err = p.unplugUSBIfConnected(ctx, s) + if err != nil { + return false, err + } + + err = p.configureAndPlugUSB(ctx, s, hashSum) + if err != nil { + return false, err + } + + s.Printf("Image mounted successfully: %s\n", filepath.Base(imagePath)) + + return true, nil +} + +// transferAndMount transfers the image from the client to the dutagent, uploads +// it to the PiKVM if missing, and mounts it. +func (p *PiKVM) transferAndMount(ctx context.Context, s module.Session, imagePath, precomputedHash string, precomputedSize int64) error { + s.Println("Requesting file from client...") + + fileReader, err := s.RequestFile(imagePath) + if err != nil { + return fmt.Errorf("failed to request file from client: %w", err) + } + + tmpFile, err := os.CreateTemp("", "pikvm-image-*") + if err != nil { + return fmt.Errorf("failed to create temporary file: %w", err) + } + + tmpPath := tmpFile.Name() + defer os.Remove(tmpPath) // Clean up temp file when done + + s.Println("Transferring file from client...") + + bytesWritten, err := io.Copy(tmpFile, fileReader) + if err != nil { + tmpFile.Close() + + return fmt.Errorf("failed to transfer file from client: %w", err) + } + + tmpFile.Close() + + s.Printf("Transferred %d bytes from client\n", bytesWritten) + + hashSum := precomputedHash + fileSize := precomputedSize + + // If hash wasn't precomputed, calculate it now. + if hashSum == "" { + hashSum, err = calcSHA256(tmpPath) + if err != nil { + return fmt.Errorf("failed to calculate image hash: %v", err) + } + + s.Printf("Image hash: %s\n", hashSum) + + fileSize = bytesWritten + } + + err = p.unplugUSBIfConnected(ctx, s) + if err != nil { + return err + } + + err = p.uploadImageIfMissing(ctx, s, tmpPath, hashSum, fileSize) + if err != nil { + return err + } + + err = p.configureAndPlugUSB(ctx, s, hashSum) + if err != nil { + return err + } + + s.Printf("Image mounted successfully: %s\n", filepath.Base(imagePath)) + + return nil +} + +// unplugUSBIfConnected checks if USB is connected and unplugs it if necessary. +func (p *PiKVM) unplugUSBIfConnected(ctx context.Context, s module.Session) error { + status, err := p.getStatus(ctx) + if err != nil { + return fmt.Errorf("failed to check USB plug state: %v", err) + } + + if !status.Result.Drive.Connected { + return nil + } + + s.Println("Unplugging USB port...") + + resp, err := p.doRequest(ctx, http.MethodPost, "/api/msd/set_connected?connected=0", nil, "", requestOptions{}) + if err != nil { + return fmt.Errorf("failed to unplug USB device: %v", err) + } + + resp.Body.Close() + + return nil +} + +// uploadImageIfMissing checks if image exists and uploads it if missing. +func (p *PiKVM) uploadImageIfMissing(ctx context.Context, s module.Session, imagePath, hashSum string, size int64) error { + err := p.checkImageExists(ctx, hashSum) + if err == nil { + s.Println("Image already exists in storage.") + + return nil + } + + if !errors.Is(err, ErrMissingImage) { + return fmt.Errorf("failed to check if image exists: %v", err) + } + + s.Println("Image not found in storage, preparing to upload.") + + // Ensure there's enough free space + err = p.ensureFreeSpace(ctx, s, size) + if err != nil { + return fmt.Errorf("failed to ensure free space: %v", err) + } + + s.Printf("Uploading image file: %s (%d bytes)\n", filepath.Base(imagePath), size) + + file, err := os.Open(imagePath) + if err != nil { + return fmt.Errorf("failed to open image file for upload: %v", err) + } + defer file.Close() + + // Upload the image file + err = p.uploadImageFile(ctx, file, hashSum, size, filepath.Base(imagePath)) + if err != nil { + return err + } + + s.Println("Image uploaded successfully.") + + return nil +} + +// uploadImageFile uploads an image file to PiKVM storage using multipart upload. +func (p *PiKVM) uploadImageFile(ctx context.Context, file *os.File, hashSum string, size int64, filename string) error { + // Build multipart header and footer without streaming the body + var buf bytes.Buffer + + writer := multipart.NewWriter(&buf) + + // Create the form file header + _, err := writer.CreateFormFile("file", filename) + if err != nil { + return fmt.Errorf("failed to create form file: %w", err) + } + + header := make([]byte, buf.Len()) + copy(header, buf.Bytes()) + + // Close writer to emit the closing boundary + err = writer.Close() + if err != nil { + return fmt.Errorf("failed to close multipart writer: %w", err) + } + + footer := buf.Bytes()[len(header):] + + contentLength := int64(len(header)) + size + int64(len(footer)) + body := io.MultiReader(bytes.NewReader(header), file, bytes.NewReader(footer)) + + uploadResp, err := p.doRequest( + ctx, + http.MethodPost, + fmt.Sprintf("/api/msd/write?image=%s", hashSum), + body, + writer.FormDataContentType(), + requestOptions{contentLength: contentLength, noTimeout: true}, + ) + if err != nil { + return fmt.Errorf("failed to upload image: %v", err) + } + defer uploadResp.Body.Close() + + return nil +} + +// configureAndPlugUSB configures the USB device with the given image and plugs it in. +// The imageName can be either a hash (for uploaded images) or a filename (for URL-mounted images). +func (p *PiKVM) configureAndPlugUSB(ctx context.Context, s module.Session, imageName string) error { + s.Println("Configuring image...") + + configEndpoint := fmt.Sprintf("/api/msd/set_params?image=%s&cdrom=0&rw=1", imageName) + + configResp, err := p.doRequest(ctx, http.MethodPost, configEndpoint, nil, "", requestOptions{}) + if err != nil { + return fmt.Errorf("failed to configure USB device: %v", err) + } + + configResp.Body.Close() + + s.Println("Plugging USB port...") + + plugResp, err := p.doRequest(ctx, http.MethodPost, "/api/msd/set_connected?connected=1", nil, "", requestOptions{}) + if err != nil { + return fmt.Errorf("failed to plug USB device: %v", err) + } + + plugResp.Body.Close() + + return nil +} + +func (p *PiKVM) handleMountURL(ctx context.Context, s module.Session, imageURL string) error { + s.Printf("Mounting image from URL: %s\n", imageURL) + + // Ensure MSD is disconnected before remote download + err := p.unplugUSBIfConnected(ctx, s) + if err != nil { + return err + } + + imageName := filepath.Base(imageURL) + + // If image already exists, skip download and just mount + exists, err := p.checkImageAlreadyExists(ctx, imageName) + if err == nil && exists { + s.Printf("Warning: image %s already exists on PiKVM, reusing without download\n", imageName) + + return p.configureAndPlugUSB(ctx, s, imageName) + } + + // Download the image from URL + err = p.downloadRemoteImage(ctx, s, imageURL, imageName) + if err != nil { + return err + } + + // Configure and connect the downloaded (or existing) image by name after it is present/complete + err = p.waitForImage(ctx, imageName, imageWaitTimeout, imageWaitInterval) + if err != nil { + return err + } + + err = p.configureAndPlugUSB(ctx, s, imageName) + if err != nil { + return err + } + + s.Printf("Image mounted successfully from URL: %s\n", imageURL) + + return nil +} + +// checkImageAlreadyExists checks if an image with the given name already exists and is complete. +func (p *PiKVM) checkImageAlreadyExists(ctx context.Context, imageName string) (bool, error) { + status, err := p.getStatus(ctx) + if err != nil { + return false, err + } + + img, ok := status.Result.Storage.Images[imageName] + + return ok && img.Complete, nil +} + +// downloadRemoteImage downloads an image from a URL to PiKVM storage. +func (p *PiKVM) downloadRemoteImage(ctx context.Context, s module.Session, imageURL, imageName string) error { + // Construct endpoint with URL encoding + endpoint := fmt.Sprintf("/api/msd/write_remote?url=%s&timeout=%d&image=%s", + url.QueryEscape(imageURL), + remoteDownloadTimout, + url.QueryEscape(imageName)) + + resp, err := p.doRequest(ctx, http.MethodPost, endpoint, nil, "", requestOptions{noTimeout: true}) + if err != nil { + return fmt.Errorf("failed to mount image from URL: %v", err) + } + defer resp.Body.Close() + + bodyBytes, err := io.ReadAll(resp.Body) + if err != nil { + return fmt.Errorf("failed to read write_remote response: %v", err) + } + + return p.handleRemoteDownloadResponse(s, imageName, bodyBytes) +} + +// remoteImageResponse represents the response from the PiKVM remote image download API. +type remoteImageResponse struct { + Ok bool `json:"ok"` + Result struct { + Error string `json:"error"` + } `json:"result"` +} + +// handleRemoteDownloadResponse parses and handles the streaming JSON response from remote download. +func (p *PiKVM) handleRemoteDownloadResponse(s module.Session, imageName string, bodyBytes []byte) error { + res, err := parseStreamingJSON(bodyBytes) + if err != nil { + return err + } + + if !res.Ok { + if res.Result.Error == "MsdImageExistsError" { + s.Printf("Warning: image %s already exists on PiKVM, reusing without download\n", imageName) + + return nil + } + + return fmt.Errorf("failed to mount image from URL: %s", string(bodyBytes)) + } + + return nil +} + +// parseStreamingJSON parses the last valid JSON line from a streaming response. +func parseStreamingJSON(bodyBytes []byte) (*remoteImageResponse, error) { + var res remoteImageResponse + + // PiKVM streams JSON progress updates line-by-line during download + // Parse the last valid JSON line as the final result + lines := bytes.Split(bodyBytes, []byte("\n")) + + // Try from the last line backwards to find valid JSON + for i := len(lines) - 1; i >= 0; i-- { + line := bytes.TrimSpace(lines[i]) + if len(line) == 0 { + continue + } + + // Try to parse this line as JSON + err := json.Unmarshal(line, &res) + if err == nil { + return &res, nil + } + } + + maxLen := 500 + if len(bodyBytes) < maxLen { + maxLen = len(bodyBytes) + } + + return nil, fmt.Errorf("no valid JSON found in write_remote response. Response (first %d bytes): %s", + maxLen, string(bodyBytes[:maxLen])) +} + +func (p *PiKVM) handleUnmount(ctx context.Context, s module.Session) error { + resp, err := p.doRequest(ctx, http.MethodPost, "/api/msd/set_connected?connected=0", nil, "", requestOptions{}) + if err != nil { + return fmt.Errorf("failed to unmount media: %v", err) + } + defer resp.Body.Close() + + s.Println("Virtual media unmounted successfully") + + return nil +} + +func (p *PiKVM) handleMediaStatus(ctx context.Context, s module.Session) error { + status, err := p.getStatus(ctx) + if err != nil { + return fmt.Errorf("failed to get media status: %v", err) + } + + connected := status.Result.Drive.Connected + + imageName := mediaNone + if connected && status.Result.Drive.Image.Name != "" { + imageName = status.Result.Drive.Image.Name + } + + s.Printf("Virtual media status:\n") + s.Printf(" Connected: %v\n", connected) + s.Printf(" Image: %s\n", imageName) + + return nil +} + +// getStatus retrieves the current MSD status from PiKVM. +func (p *PiKVM) getStatus(ctx context.Context) (*Status, error) { + resp, err := p.doRequest(ctx, http.MethodGet, "/api/msd", nil, "", requestOptions{}) + if err != nil { + return nil, fmt.Errorf("failed to get MSD status: %w", err) + } + defer resp.Body.Close() + + var status Status + + err = json.NewDecoder(resp.Body).Decode(&status) + if err != nil { + return nil, fmt.Errorf("failed to decode status: %w", err) + } + + if !status.Ok { + return nil, fmt.Errorf("status response not ok") + } + + if status.Result.Busy { + return nil, fmt.Errorf("PiKVM mass-storage is busy") + } + + if !status.Result.Enabled || !status.Result.Online { + return nil, fmt.Errorf("PiKVM mass-storage is not enabled or online") + } + + return &status, nil +} + +// calcSHA256 calculates the SHA256 hash of a file. +func calcSHA256(path string) (string, error) { + file, err := os.Open(path) + if err != nil { + return "", fmt.Errorf("failed to open file: %w", err) + } + defer file.Close() + + hash := sha256.New() + + _, err = io.Copy(hash, file) + if err != nil { + return "", fmt.Errorf("failed to calculate SHA256: %w", err) + } + + return fmt.Sprintf("%x", hash.Sum(nil)), nil +} + +// checkImageExists checks if an image with the given hash already exists in storage. +func (p *PiKVM) checkImageExists(ctx context.Context, hashSum string) error { + status, err := p.getStatus(ctx) + if err != nil { + return err + } + + for imageHash := range status.Result.Storage.Images { + if imageHash == hashSum { + return nil + } + } + + return ErrMissingImage +} + +// getFreeSpace returns the total free space across all writable storage partitions. +func getFreeSpace(status *Status) (uint64, error) { + if len(status.Result.Storage.Parts) == 0 { + return 0, fmt.Errorf("no storage partitions available") + } + + var totalFree uint64 + + foundWritable := false + + for _, part := range status.Result.Storage.Parts { + if part.Writable { + totalFree += part.Free + foundWritable = true + } + } + + if !foundWritable { + return 0, fmt.Errorf("no writable storage partition found") + } + + return totalFree, nil +} + +// deleteImage deletes an image from PiKVM storage by its name (hash). +func (p *PiKVM) deleteImage(ctx context.Context, imageName string) error { + resp, err := p.doRequest(ctx, http.MethodPost, fmt.Sprintf("/api/msd/remove?image=%s", imageName), nil, "", requestOptions{}) + if err != nil { + return fmt.Errorf("failed to delete image: %w", err) + } + defer resp.Body.Close() + + var response struct { + Ok bool `json:"ok"` + } + + err = json.NewDecoder(resp.Body).Decode(&response) + if err != nil { + return fmt.Errorf("failed to decode delete response: %w", err) + } + + if !response.Ok { + return fmt.Errorf("delete operation failed") + } + + return nil +} + +// findOldestImage finds the image with the oldest modification timestamp. +// If skipImage is not empty, that image will be excluded from consideration. +func findOldestImage(status *Status, skipImage string) (string, error) { + if len(status.Result.Storage.Images) == 0 { + return "", fmt.Errorf("no images available to delete") + } + + var oldestName string + + var oldestTime float32 + + var found bool + + for name, img := range status.Result.Storage.Images { + // Skip the image we want to exclude + if name == skipImage { + continue + } + + if !found || img.ModTS < oldestTime { + oldestTime = img.ModTS + oldestName = name + found = true + } + } + + if oldestName == "" { + return "", fmt.Errorf("no deletable images available (only the currently connected image exists)") + } + + return oldestName, nil +} + +// ensureFreeSpace checks if there's enough free space for the given size. +// If not, it keeps deleting the oldest images until there's enough space. +// Currently connected image (if any) is not deleted. +func (p *PiKVM) ensureFreeSpace(ctx context.Context, s module.Session, requiredSize int64) error { + // Validate input + if requiredSize < 0 { + return fmt.Errorf("invalid required size: %d", requiredSize) + } + + deletionCount := 0 + + requiredSizeUint := uint64(requiredSize) // #nosec G115 -- validated above + + // Keep deleting oldest images until we have enough space + for { + status, err := p.getStatus(ctx) + if err != nil { + return err + } + + freeSpace, err := getFreeSpace(status) + if err != nil { + return fmt.Errorf("failed to get free space: %w", err) + } + + // Check if we have enough space + if requiredSizeUint+safetyMargin <= freeSpace { + if deletionCount > 0 { + s.Printf("Freed sufficient space by deleting %d old image(s).\n", deletionCount) + } + + break + } + + // Safety check: prevent infinite loop + if deletionCount >= maxDeletionRetries { + return fmt.Errorf( + "insufficient storage space after attempting to delete %d image(s): need %d bytes, have %d bytes free", + deletionCount, requiredSize, freeSpace, + ) + } + + // Delete one old image and continue + err = p.deleteOldImage(ctx, s, status, requiredSizeUint, freeSpace) + if err != nil { + return err + } + + deletionCount++ + + time.Sleep(scanningDelay) // wait a bit for the storage to update + } + + return nil +} + +// deleteOldImage finds and deletes the oldest image to free up space. +func (p *PiKVM) deleteOldImage(ctx context.Context, s module.Session, status *Status, requiredSize, freeSpace uint64) error { + // Determine which image to skip (if currently connected) + var skipImage string + if status.Result.Drive.Connected { + skipImage = status.Result.Drive.Image.Name + } + + // Find the oldest image to delete + oldestImage, err := findOldestImage(status, skipImage) + if err != nil { + return fmt.Errorf("insufficient storage space: need %d bytes, have %d bytes free: %w", requiredSize, freeSpace, err) + } + + s.Printf("Deleting old image to free space (need %d MB, have %d MB free)...\n", + (requiredSize+safetyMargin)/bytesPerMB, freeSpace/bytesPerMB) + + err = p.deleteImage(ctx, oldestImage) + + return err +} + +// waitForImage polls MSD status until the image is present and complete or times out. +func (p *PiKVM) waitForImage(ctx context.Context, imageName string, timeout, interval time.Duration) error { + deadline := time.Now().Add(timeout) + + for time.Now().Before(deadline) { + status, err := p.getStatus(ctx) + if err != nil { + return fmt.Errorf("failed to get status while waiting for image: %w", err) + } + + if img, ok := status.Result.Storage.Images[imageName]; ok && img.Complete { + return nil + } + + time.Sleep(interval) + } + + return fmt.Errorf("image %s not ready after %v", imageName, timeout) +} diff --git a/pkg/module/pikvm/pikvm-example-cfg.yml b/pkg/module/pikvm/pikvm-example-cfg.yml new file mode 100644 index 0000000..8a2f6af --- /dev/null +++ b/pkg/module/pikvm/pikvm-example-cfg.yml @@ -0,0 +1,53 @@ +version: 0 +devices: + my-server: + desc: "A server controlled via PiKVM" + cmds: + # Power Management + power: + desc: "Power management: on|off|force-off|reset|force-reset|status" + modules: + - module: pikvm + main: true + options: + host: https://pikvm.local + user: admin + password: admin + timeout: 10s + mode: power + + # Keyboard Control + keyboard: + desc: "Keyboard control: type |key |key-combo " + modules: + - module: pikvm + main: true + options: + host: https://pikvm.local + user: admin + password: admin + mode: keyboard + + # Virtual Media + media: + desc: "Virtual media control: mount |mount-url |unmount|media-status" + modules: + - module: pikvm + main: true + options: + host: https://pikvm.local + user: admin + password: admin + mode: media + + # Screenshot + screenshot: + desc: "Capture a screenshot from PiKVM" + modules: + - module: pikvm + main: true + options: + host: https://pikvm.local + user: admin + password: admin + mode: screenshot diff --git a/pkg/module/pikvm/pikvm-example-keyboard.yml b/pkg/module/pikvm/pikvm-example-keyboard.yml new file mode 100644 index 0000000..b143f0b --- /dev/null +++ b/pkg/module/pikvm/pikvm-example-keyboard.yml @@ -0,0 +1,17 @@ +version: 0 +devices: + my-server: + desc: "Server controlled via PiKVM" + cmds: + keyboard: + desc: "Keyboard control: type |key |key-combo " + modules: + - module: pikvm + main: true + options: + host: https://pikvm.local + user: admin + password: admin + mode: keyboard +# Usage: dutctl my-server keyboard key F12 +# Usage (custom key delay): dutctl my-server keyboard --delay 50ms key-combo Ctrl+Alt+Delete diff --git a/pkg/module/pikvm/pikvm-example-media.yml b/pkg/module/pikvm/pikvm-example-media.yml new file mode 100644 index 0000000..b4872f3 --- /dev/null +++ b/pkg/module/pikvm/pikvm-example-media.yml @@ -0,0 +1,16 @@ +version: 0 +devices: + my-server: + desc: "Server controlled via PiKVM" + cmds: + media: + desc: "Virtual media control: mount |mount-url |unmount|media-status" + modules: + - module: pikvm + main: true + options: + host: https://pikvm.local + user: admin + password: admin + mode: media +# Usage: dutctl my-server media mount-url https://releases.ubuntu.com/22.04/ubuntu-22.04-live-server-amd64.iso diff --git a/pkg/module/pikvm/pikvm-example-power.yml b/pkg/module/pikvm/pikvm-example-power.yml new file mode 100644 index 0000000..568181b --- /dev/null +++ b/pkg/module/pikvm/pikvm-example-power.yml @@ -0,0 +1,16 @@ +version: 0 +devices: + my-server: + desc: "Server controlled via PiKVM" + cmds: + power: + desc: "Power management: on|off|force-off|reset|force-reset|status" + modules: + - module: pikvm + main: true + options: + host: https://pikvm.local + user: admin + password: admin + mode: power +# Usage: dutctl my-server power status diff --git a/pkg/module/pikvm/pikvm-example-screenshot.yml b/pkg/module/pikvm/pikvm-example-screenshot.yml new file mode 100644 index 0000000..2a47c28 --- /dev/null +++ b/pkg/module/pikvm/pikvm-example-screenshot.yml @@ -0,0 +1,17 @@ +version: 0 +devices: + my-server: + desc: "Server controlled via PiKVM" + cmds: + screenshot: + desc: "Capture a screenshot from PiKVM" + modules: + - module: pikvm + main: true + options: + host: https://pikvm.local + user: admin + password: admin + mode: screenshot +# Usage: dutctl my-server screenshot +# The screenshot will be saved to the current directory diff --git a/pkg/module/pikvm/pikvm-key-strings.md b/pkg/module/pikvm/pikvm-key-strings.md new file mode 100644 index 0000000..77547c1 --- /dev/null +++ b/pkg/module/pikvm/pikvm-key-strings.md @@ -0,0 +1,62 @@ +# PiKVM keyboard key strings + +This module accepts the following key strings for `keyboard key ` and for `keyboard key-combo ` (e.g. `Ctrl+Alt+Delete`). + +- Key matching is case-insensitive. +- The list below is the authoritative list from `pkg/module/pikvm/keyboard.go` (`keyboardMap`). + +## Letters +`a` `b` `c` `d` `e` `f` `g` `h` `i` `j` `k` `l` `m` `n` `o` `p` `q` `r` `s` `t` `u` `v` `w` `x` `y` `z` + +## Digits +`0` `1` `2` `3` `4` `5` `6` `7` `8` `9` + +## Modifiers +`ctrl` `control` `alt` `shift` `meta` `win` `windows` `cmd` `command` + +## Common keys +`space` +`enter` `return` +`tab` +`esc` `escape` +`backspace` +`delete` `del` +`insert` `ins` +`home` +`end` +`pageup` `pgup` +`pagedown` `pgdn` + +## Function keys +`f1` `f2` `f3` `f4` `f5` `f6` `f7` `f8` `f9` `f10` `f11` `f12` + +## Arrow keys +`up` `arrowup` +`down` `arrowdown` +`left` `arrowleft` +`right` `arrowright` + +## Punctuation and symbols +`-` `minus` +`=` `equals` +`[` `leftbracket` +`]` `rightbracket` +`\\` `backslash` +`;` `semicolon` +`'` `quote` +`` ` `` `backquote` +`,` `comma` +`.` `period` +`/` `slash` + +## Additional keys +`capslock` +`printscreen` +`pause` +`scrolllock` `scroll` +`numlock` `num` +`contextmenu` `menu` + +## Numpad keys +`numpad0` `numpad1` `numpad2` `numpad3` `numpad4` `numpad5` `numpad6` `numpad7` `numpad8` `numpad9` +`numpad/` `numpad*` `numpad-` `numpad+` `numpadenter` `numpad.` diff --git a/pkg/module/pikvm/pikvm.go b/pkg/module/pikvm/pikvm.go new file mode 100644 index 0000000..ad5fe69 --- /dev/null +++ b/pkg/module/pikvm/pikvm.go @@ -0,0 +1,340 @@ +// Copyright 2025 Blindspot Software +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +// Package pikvm provides a dutagent module that allows control of a PiKVM device via HTTP API. +package pikvm + +import ( + "context" + "crypto/tls" + "fmt" + "io" + "log" + "net/http" + "net/url" + "path" + "strings" + "time" + + "github.com/BlindspotSoftware/dutctl/pkg/module" +) + +func init() { + module.Register(module.Record{ + ID: "pikvm", + New: func() module.Module { + return &PiKVM{} + }, + }) +} + +// PiKVM is a module that provides power management, keyboard control, and virtual media +// functionality for a DUT via PiKVM device. +type PiKVM struct { + Host string // Host is the address of the PiKVM device (e.g., "192.168.1.100" or "https://pikvm.local") + User string // User for authentication (default: "admin") + Password string // Password for authentication + Timeout string // Timeout for HTTP requests (default: "10s") + Mode string // Required: Mode to execute (must be: "power", "keyboard", "media", or "screenshot") + + client *http.Client + baseURL *url.URL +} + +// Ensure implementing the Module interface. +var _ module.Module = &PiKVM{} + +const ( + defaultUser = "admin" + defaultTimeout = 10 * time.Second + minArgsRequired = 2 // Minimum arguments required for commands with parameters +) + +// Mode constants. +const ( + cmdTypePower = "power" + cmdTypeKeyboard = "keyboard" + cmdTypeMedia = "media" + cmdTypeScreenshot = "screenshot" +) + +func (p *PiKVM) Help() string { + log.Println("pikvm module: Help called") + + help := strings.Builder{} + help.WriteString("PiKVM Control Module\n") + help.WriteString("PiKVM is a Raspberry Pi based KVM-over-IP device (KVM = Keyboard, Video, Mouse).\n") + help.WriteString("This module provides comprehensive control of a PiKVM device.\n\n") + + switch p.Mode { + case cmdTypePower: + writePowerHelp(&help) + case cmdTypeKeyboard: + writeKeyboardHelp(&help) + case cmdTypeMedia: + writeMediaHelp(&help) + case cmdTypeScreenshot: + writeScreenshotHelp(&help) + default: + writePowerHelp(&help) + writeKeyboardHelp(&help) + writeMediaHelp(&help) + writeScreenshotHelp(&help) + } + + help.WriteString("Configured PiKVM: " + p.Host + "\n") + + return help.String() +} + +func writePowerHelp(help *strings.Builder) { + help.WriteString("Power Management:\n") + help.WriteString(" pikvm on - Power on (does nothing if already on)\n") + help.WriteString(" pikvm off - Graceful shutdown (soft power-off)\n") + help.WriteString(" pikvm force-off - Force power off (hard shutdown, 5+ second press)\n") + help.WriteString(" pikvm reset - Reset via ATX reset button\n") + help.WriteString(" pikvm force-reset - Force reset (hardware hot reset)\n") + help.WriteString(" pikvm status - Query current power state\n\n") +} + +func writeKeyboardHelp(help *strings.Builder) { + help.WriteString("Keyboard Control:\n") + help.WriteString(" pikvm [--delay ] ... - Optional key delay (default: 500ms)\n") + help.WriteString(" pikvm type - Type a text string\n") + help.WriteString(" pikvm key - Send a single key (e.g., Enter, Escape, F12)\n") + help.WriteString(" pikvm key-combo - Send key combination (e.g., Ctrl+Alt+Delete)\n\n") +} + +func writeMediaHelp(help *strings.Builder) { + help.WriteString("Virtual Media:\n") + help.WriteString(" pikvm mount - Mount an image file from local filesystem\n") + help.WriteString(" pikvm mount-url - Mount an image from a URL\n") + help.WriteString(" pikvm unmount - Unmount current virtual media\n") + help.WriteString(" pikvm media-status - Show mounted media information\n\n") +} + +func writeScreenshotHelp(help *strings.Builder) { + help.WriteString("Screenshot:\n") + help.WriteString(" pikvm screenshot - Capture a screenshot (saved to current directory)\n\n") +} + +func (p *PiKVM) Init() error { + log.Printf("pikvm module: Init starting for host %s", p.Host) + + err := p.validateMode() + if err != nil { + return err + } + + // Set default user if not provided + if p.User == "" { + p.User = defaultUser + log.Printf("pikvm module: Using default user '%s'", defaultUser) + } + + // Parse custom timeout if provided + timeout := defaultTimeout + + if p.Timeout != "" { + parsedTimeout, err := time.ParseDuration(p.Timeout) + if err == nil { + timeout = parsedTimeout + log.Printf("pikvm module: Using custom timeout %v", timeout) + } else { + log.Printf("pikvm module: Invalid timeout format '%s', using default %v", p.Timeout, defaultTimeout) + } + } + + // Create HTTP client with TLS config that accepts self-signed certificates + // This is necessary because PiKVM devices typically use self-signed certs + transport := &http.Transport{ + TLSClientConfig: &tls.Config{ + InsecureSkipVerify: true, //nolint:gosec // PiKVM devices use self-signed certificates + }, + } + p.client = &http.Client{ + Timeout: timeout, + Transport: transport, + } + + log.Printf("pikvm module: Init completed successfully for %s", p.baseURL.String()) + + return nil +} + +func (p *PiKVM) Deinit() error { + log.Println("pikvm module: Deinit called") + + return nil +} + +func (p *PiKVM) Run(ctx context.Context, s module.Session, args ...string) error { + if p.client == nil { + return fmt.Errorf("pikvm: client not initialized") + } + + // Route based on the configured mode + switch p.Mode { + case cmdTypePower: + return p.handlePowerCommandRouter(ctx, s, args) + case cmdTypeKeyboard: + return p.handleKeyboardCommandRouter(ctx, s, args) + case cmdTypeMedia: + return p.handleMediaCommandRouter(ctx, s, args) + case cmdTypeScreenshot: + return p.handleScreenshot(ctx, s) + default: + return fmt.Errorf("pikvm: invalid configuration: unknown mode %q", p.Mode) + } +} + +func (p *PiKVM) validateMode() error { + baseURL, err := normalizeAndParseHost(p.Host) + if err != nil { + return err + } + // store normalized/parsed host for later request building + p.baseURL = baseURL + + if p.Mode == "" { + return fmt.Errorf("pikvm: mode option is required (must be: %s, %s, %s, or %s)", + cmdTypePower, cmdTypeKeyboard, cmdTypeMedia, cmdTypeScreenshot) + } + + switch p.Mode { + case cmdTypePower, cmdTypeKeyboard, cmdTypeMedia, cmdTypeScreenshot: + return nil + default: + return fmt.Errorf("pikvm: invalid mode %q (must be: %s, %s, %s, or %s)", + p.Mode, cmdTypePower, cmdTypeKeyboard, cmdTypeMedia, cmdTypeScreenshot) + } +} + +type requestOptions struct { + contentLength int64 + noTimeout bool +} + +// doRequest is the core request method that handles all HTTP requests. +func (p *PiKVM) doRequest( + ctx context.Context, + method, urlPath string, + body io.Reader, + contentType string, + opts requestOptions, +) (*http.Response, error) { + // Build URL + targetURL, err := p.buildRequestURL(urlPath) + if err != nil { + return nil, err + } + + // Create request + req, err := http.NewRequestWithContext(ctx, method, targetURL, body) + if err != nil { + return nil, err + } + + // Set headers + p.setRequestHeaders(req, contentType, opts) + + // Choose client (with or without timeout) + client := p.selectHTTPClient(opts) + + // Execute request + resp, err := client.Do(req) + if err != nil { + return nil, err + } + + // Check status + err = checkResponseStatus(resp) + if err != nil { + return nil, err + } + + return resp, nil +} + +// buildRequestURL constructs the full request URL from base URL and path. +func (p *PiKVM) buildRequestURL(urlPath string) (string, error) { + if p.baseURL == nil { + return "", fmt.Errorf("pikvm: base URL is not initialized") + } + + targetURL := *p.baseURL + + pathPart := urlPath + rawQuery := "" + + if before, after, ok := strings.Cut(urlPath, "?"); ok { + pathPart = before + rawQuery = after + } + + targetURL.Path = path.Join(targetURL.Path, pathPart) + targetURL.RawQuery = rawQuery + + return targetURL.String(), nil +} + +// setRequestHeaders sets authentication and content type headers on the request. +func (p *PiKVM) setRequestHeaders(req *http.Request, contentType string, opts requestOptions) { + if p.User != "" && p.Password != "" { + req.SetBasicAuth(p.User, p.Password) + } + + if contentType != "" { + req.Header.Set("Content-Type", contentType) + } + + if opts.contentLength > 0 { + req.ContentLength = opts.contentLength + } +} + +// selectHTTPClient returns the appropriate HTTP client based on options. +func (p *PiKVM) selectHTTPClient(opts requestOptions) *http.Client { + if opts.noTimeout { + return &http.Client{ + Transport: &http.Transport{ + TLSClientConfig: &tls.Config{ + InsecureSkipVerify: true, //nolint:gosec // PiKVM uses self-signed certificates + }, + }, + } + } + + return p.client +} + +// checkResponseStatus validates the HTTP response status code. +func checkResponseStatus(resp *http.Response) error { + if resp.StatusCode >= 200 && resp.StatusCode < 300 { + return nil + } + + bodyBytes, _ := io.ReadAll(resp.Body) + resp.Body.Close() + + return fmt.Errorf("pikvm: API returned %s: %s", resp.Status, string(bodyBytes)) +} + +func normalizeAndParseHost(host string) (*url.URL, error) { + host = strings.TrimSpace(host) + if host == "" { + return nil, fmt.Errorf("pikvm: host is not set") + } + + if !strings.HasPrefix(host, "http://") && !strings.HasPrefix(host, "https://") { + host = "https://" + host + } + + baseURL, err := url.Parse(strings.TrimRight(host, "/")) + if err != nil { + return nil, fmt.Errorf("pikvm: invalid host URL: %v", err) + } + + return baseURL, nil +} diff --git a/pkg/module/pikvm/pikvm_test.go b/pkg/module/pikvm/pikvm_test.go new file mode 100644 index 0000000..971d376 --- /dev/null +++ b/pkg/module/pikvm/pikvm_test.go @@ -0,0 +1,250 @@ +// Copyright 2025 Blindspot Software +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package pikvm + +import ( + "strings" + "testing" +) + +func TestInitWithDefaults(t *testing.T) { + p := &PiKVM{ + Host: "pikvm.local", + Password: "admin", + Mode: "power", + } + + err := p.Init() + if err != nil { + t.Fatalf("Init failed: %v", err) + } + + // Check default user is set + if p.User != defaultUser { + t.Fatalf("expected default user %q, got %q", defaultUser, p.User) + } + + // Check baseURL is properly set + if p.baseURL == nil { + t.Fatalf("baseURL is nil") + } + + // Check HTTPS is added if no scheme provided + if p.baseURL.Scheme != "https" { + t.Fatalf("expected https scheme, got %q", p.baseURL.Scheme) + } +} + +func TestInitWithExplicitHTTP(t *testing.T) { + p := &PiKVM{ + Host: "http://192.168.1.100", + User: "admin", + Password: "admin", + Mode: "keyboard", + } + + err := p.Init() + if err != nil { + t.Fatalf("Init failed: %v", err) + } + + if p.baseURL.Scheme != "http" { + t.Fatalf("expected http scheme, got %q", p.baseURL.Scheme) + } +} + +func TestInitWithCustomTimeout(t *testing.T) { + p := &PiKVM{ + Host: "pikvm.local", + User: "admin", + Password: "admin", + Timeout: "30s", + Mode: "media", + } + + err := p.Init() + if err != nil { + t.Fatalf("Init failed: %v", err) + } + + // Just verify it doesn't fail - actual timeout value is internal to client +} + +func TestInitWithInvalidTimeout(t *testing.T) { + p := &PiKVM{ + Host: "pikvm.local", + User: "admin", + Password: "admin", + Timeout: "invalid", + Mode: "screenshot", + } + + // Should still succeed but fall back to default timeout + err := p.Init() + if err != nil { + t.Fatalf("Init should succeed with invalid timeout (falls back to default): %v", err) + } +} + +func TestInitWithoutHost(t *testing.T) { + p := &PiKVM{ + User: "admin", + Password: "admin", + } + + err := p.Init() + if err == nil { + t.Fatalf("Init should fail when Host is not set") + } + + if !strings.Contains(err.Error(), "host is not set") { + t.Fatalf("expected error about host not set, got: %v", err) + } +} + +func TestHelp(t *testing.T) { + p := &PiKVM{ + Host: "pikvm.local", + } + + help := p.Help() + if help == "" { + t.Fatalf("Help returned empty string") + } + + // Check that help contains key sections + expectedSections := []string{ + "PiKVM Control Module", + "Power Management", + "Keyboard Control", + "Virtual Media", + "on", + "off", + "force-off", + "reset", + "type", + "key", + "key-combo", + "mount", + "unmount", + } + + for _, section := range expectedSections { + if !strings.Contains(help, section) { + t.Fatalf("Help should contain %q, but doesn't", section) + } + } +} + +func TestHelpWithMode(t *testing.T) { + tests := []struct { + name string + mode string + mustContain []string + mustNotContain []string + }{ + { + name: "power", + mode: cmdTypePower, + mustContain: []string{ + "Power Management", + "pikvm on", + "Configured PiKVM: pikvm.local", + }, + mustNotContain: []string{ + "Keyboard Control", + "Virtual Media", + "Screenshot:", + }, + }, + { + name: "keyboard", + mode: cmdTypeKeyboard, + mustContain: []string{ + "Keyboard Control", + "--delay", + "pikvm type", + "Configured PiKVM: pikvm.local", + }, + mustNotContain: []string{ + "Power Management", + "Virtual Media", + "Screenshot:", + }, + }, + { + name: "media", + mode: cmdTypeMedia, + mustContain: []string{ + "Virtual Media", + "pikvm mount", + "Configured PiKVM: pikvm.local", + }, + mustNotContain: []string{ + "Power Management", + "Keyboard Control", + "Screenshot:", + }, + }, + { + name: "screenshot", + mode: cmdTypeScreenshot, + mustContain: []string{ + "Screenshot:", + "pikvm screenshot", + "Configured PiKVM: pikvm.local", + }, + mustNotContain: []string{ + "Power Management", + "Keyboard Control", + "Virtual Media", + }, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + p := &PiKVM{Host: "pikvm.local", Mode: tc.mode} + help := p.Help() + for _, s := range tc.mustContain { + if !strings.Contains(help, s) { + t.Fatalf("Help should contain %q, but doesn't", s) + } + } + for _, s := range tc.mustNotContain { + if strings.Contains(help, s) { + t.Fatalf("Help should not contain %q, but does", s) + } + } + }) + } +} + +func TestDeinit(t *testing.T) { + p := &PiKVM{} + + err := p.Deinit() + if err != nil { + t.Fatalf("Deinit should not return error: %v", err) + } +} + +func TestParseKeyCombo(t *testing.T) { + // Test key combination parsing + comboStr := "Ctrl+Alt+Delete" + keys := strings.Split(comboStr, "+") + + expectedKeys := []string{"Ctrl", "Alt", "Delete"} + if len(keys) != len(expectedKeys) { + t.Fatalf("expected %d keys, got %d", len(expectedKeys), len(keys)) + } + + for i, key := range keys { + trimmed := strings.TrimSpace(key) + if trimmed != expectedKeys[i] { + t.Fatalf("expected key %q at position %d, got %q", expectedKeys[i], i, trimmed) + } + } +} diff --git a/pkg/module/pikvm/power.go b/pkg/module/pikvm/power.go new file mode 100644 index 0000000..e531b0c --- /dev/null +++ b/pkg/module/pikvm/power.go @@ -0,0 +1,137 @@ +// Copyright 2025 Blindspot Software +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package pikvm + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "strings" + + "github.com/BlindspotSoftware/dutctl/pkg/module" +) + +// Power command constants. +const ( + on = "on" + off = "off" + forceOff = "force-off" + reset = "reset" + forceReset = "force-reset" + status = "status" + + statusOn = "On" + statusOff = "Off" +) + +// handlePowerCommandRouter routes power commands based on the first argument. +func (p *PiKVM) handlePowerCommandRouter(ctx context.Context, s module.Session, args []string) error { + if len(args) == 0 { + s.Println("Power command requires an action: on|off|force-off|reset|force-reset|status") + + return nil + } + + command := strings.ToLower(args[0]) + + return p.handlePowerCommand(ctx, s, command) +} + +// handlePowerCommand dispatches power management commands. +func (p *PiKVM) handlePowerCommand(ctx context.Context, s module.Session, command string) error { + switch command { + case on: + return p.sendATXPower(ctx, s, "on") + case off: + return p.sendATXPower(ctx, s, "off") + case forceOff: + return p.sendATXPower(ctx, s, "off_hard") + case reset: + return p.sendATXClick(ctx, s, "reset") + case forceReset: + return p.sendATXPower(ctx, s, "reset_hard") + case status: + return p.handleStatusCommand(ctx, s) + default: + return fmt.Errorf("unknown power action: %s (must be: on, off, force-off, reset, force-reset, status)", command) + } +} + +// sendATXPower sends an ATX power action using the /api/atx/power endpoint. +// This endpoint is intelligent: 'on' does nothing if already on, 'off' is graceful shutdown. +func (p *PiKVM) sendATXPower(ctx context.Context, s module.Session, action string) error { + endpoint := fmt.Sprintf("/api/atx/power?action=%s", action) + + resp, err := p.doRequest(ctx, http.MethodPost, endpoint, nil, "", requestOptions{}) + if err != nil { + return fmt.Errorf("ATX power %s failed: %v", action, err) + } + defer resp.Body.Close() + + s.Printf("ATX power action '%s' completed\n", action) + + return nil +} + +// sendATXClick sends an ATX button click using the /api/atx/click endpoint. +func (p *PiKVM) sendATXClick(ctx context.Context, s module.Session, button string) error { + endpoint := fmt.Sprintf("/api/atx/click?button=%s", button) + + resp, err := p.doRequest(ctx, http.MethodPost, endpoint, nil, "", requestOptions{}) + if err != nil { + return fmt.Errorf("ATX %s button click failed: %v", button, err) + } + defer resp.Body.Close() + + s.Printf("ATX %s button clicked\n", button) + + return nil +} + +// ATXStatus represents the PiKVM ATX status response. +type ATXStatus struct { + Ok bool `json:"ok"` + Result ATXStatusResult `json:"result"` +} + +// ATXStatusResult contains the actual ATX status data. +type ATXStatusResult struct { + Enabled bool `json:"enabled"` + Leds struct { + Power bool `json:"power"` + HDD bool `json:"hdd"` + } `json:"leds"` + Busy bool `json:"busy"` +} + +func (p *PiKVM) handleStatusCommand(ctx context.Context, s module.Session) error { + resp, err := p.doRequest(ctx, http.MethodGet, "/api/atx", nil, "", requestOptions{}) + if err != nil { + return fmt.Errorf("failed to get ATX status: %v", err) + } + defer resp.Body.Close() + + var status ATXStatus + + err = json.NewDecoder(resp.Body).Decode(&status) + if err != nil { + return fmt.Errorf("failed to parse status response: %v", err) + } + + if !status.Ok { + return fmt.Errorf("ATX status response not ok") + } + + // Extract power state from response + powerState := statusOff + if status.Result.Leds.Power { + powerState = statusOn + } + + s.Printf("Device power status: %s\n", powerState) + + return nil +} diff --git a/pkg/module/pikvm/screenshot.go b/pkg/module/pikvm/screenshot.go new file mode 100644 index 0000000..e2e36d4 --- /dev/null +++ b/pkg/module/pikvm/screenshot.go @@ -0,0 +1,56 @@ +// Copyright 2025 Blindspot Software +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package pikvm + +import ( + "bytes" + "context" + "fmt" + "io" + "net/http" + "strings" + "time" + + "github.com/BlindspotSoftware/dutctl/pkg/module" +) + +// handleScreenshot captures a screenshot from the PiKVM and sends it to the client. +func (p *PiKVM) handleScreenshot(ctx context.Context, s module.Session) error { + s.Println("Capturing screenshot...") + + resp, err := p.doRequest(ctx, http.MethodGet, "/api/streamer/snapshot", nil, "", requestOptions{}) + if err != nil { + // Check if error contains "Service Unavailable" which indicates the device is likely off + if strings.Contains(err.Error(), "Service Unavailable") || strings.Contains(err.Error(), "503") { + return fmt.Errorf( + "failed to take screenshot: device appears to be powered off or video stream is unavailable (503)", + ) + } + + return fmt.Errorf("failed to capture screenshot: %v", err) + } + defer resp.Body.Close() + + // Read the screenshot data + screenshotData, err := io.ReadAll(resp.Body) + if err != nil { + return fmt.Errorf("failed to read screenshot data: %v", err) + } + + // Generate filename with timestamp + filename := fmt.Sprintf("screenshot_%s.jpg", time.Now().Format("20060102_150405")) + + s.Printf("Sending screenshot to client: %s (%d bytes)\n", filename, len(screenshotData)) + + // Send the screenshot to the client + err = s.SendFile(filename, bytes.NewReader(screenshotData)) + if err != nil { + return fmt.Errorf("failed to send screenshot to client: %v", err) + } + + s.Printf("Screenshot sent successfully: %s\n", filename) + + return nil +}