Vectrum
Sign In
147 lines · 4.8 KB · Go
package attach

import (
	"context"
	"encoding/hex"
	"errors"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"time"
)

// Reading an image off the clipboard is done by asking whichever helper the
// desktop ships, because there is no portable syscall for it and the helpers
// are what every other program uses too. Several are tried in turn: a machine
// may have both Wayland and X11 tools installed with only one of them working,
// and a helper that is missing looks exactly like a helper that found nothing.
//
// Each runs under a short deadline. These programs block when the clipboard
// owner is slow or gone — a closed window whose selection nobody took over
// hangs xclip indefinitely — and a hung paste would take the whole UI with it.

const grabTimeout = 3 * time.Second

// grabber is one way to ask for the clipboard's image.
type grabber struct {
	cmd  string
	args []string
	// decode converts the helper's stdout into image bytes. Nil means the
	// bytes are already the image.
	decode func([]byte) ([]byte, error)
}

// grabbers is the list for this machine, in the order they are tried. It is a
// variable so a test can stand in for the desktop.
var grabbers = defaultGrabbers()

func defaultGrabbers() []grabber {
	switch runtime.GOOS {
	case "darwin":
		return []grabber{
			// pngpaste is not standard, but it is the one people install, and
			// it hands back the bytes with nothing in the way.
			{cmd: "pngpaste", args: []string{"-"}},
			// AppleScript is always there. It answers with «data PNGf89504…»,
			// hex inside guillemets, which decodeAppleScript unwraps.
			{cmd: "osascript", args: []string{"-e", "the clipboard as «class PNGf»"}, decode: decodeAppleScript},
		}
	case "windows":
		return []grabber{{cmd: "powershell", args: []string{"-NoProfile", "-STA", "-Command", psScript}, decode: decodeHexLine}}
	default:
		// xsel is not here on purpose: it cannot ask for a MIME target, so it
		// would answer an image request with the clipboard's text.
		return []grabber{
			{cmd: "wl-paste", args: []string{"--no-newline", "--type", "image/png"}},
			{cmd: "xclip", args: []string{"-selection", "clipboard", "-t", "image/png", "-o"}},
		}
	}
}

// psScript pulls the clipboard image through PowerShell and writes it as hex
// on one line, because a pipe out of PowerShell mangles raw bytes.
const psScript = `Add-Type -AssemblyName System.Windows.Forms;` +
	`$i=[Windows.Forms.Clipboard]::GetImage();` +
	`if($i -eq $null){exit 1};` +
	`$m=New-Object System.IO.MemoryStream;` +
	`$i.Save($m,[System.Drawing.Imaging.ImageFormat]::Png);` +
	`[BitConverter]::ToString($m.ToArray()).Replace('-','')`

func decodeHexLine(out []byte) ([]byte, error) {
	return hex.DecodeString(strings.TrimSpace(string(out)))
}

// decodeAppleScript unwraps «data PNGf89504E47…». The guillemets are UTF-8 in
// osascript's output; what matters is the hex between the tag and the closing
// mark.
func decodeAppleScript(out []byte) ([]byte, error) {
	s := strings.TrimSpace(string(out))
	i := strings.Index(s, "PNGf")
	if i < 0 {
		return nil, ErrNoImage
	}
	s = s[i+4:]
	if j := strings.IndexAny(s, "»\""); j >= 0 {
		s = s[:j]
	}
	return hex.DecodeString(strings.TrimSpace(s))
}

// Clipboard returns the picture on the system clipboard. ErrNoImage means
// there is none — the usual case, and not worth reporting as a failure.
func Clipboard() (Image, error) {
	var missing int
	for _, g := range grabbers {
		if _, err := exec.LookPath(g.cmd); err != nil {
			missing++
			continue
		}
		out, err := run(g)
		if err != nil || len(out) == 0 {
			continue // no image in that flavour, or the helper said no
		}
		if g.decode != nil {
			if out, err = g.decode(out); err != nil || len(out) == 0 {
				continue
			}
		}
		im, err := New(out, "clipboard")
		if err == nil {
			return im, nil
		}
		if errors.Is(err, ErrTooBig) {
			// A picture was found and it is not going to fit. Say so, rather
			// than reporting an empty clipboard to someone who can see their
			// screenshot sitting on it.
			return Image{}, err
		}
	}
	if missing == len(grabbers) && len(grabbers) > 0 {
		return Image{}, errors.New("no clipboard helper installed — " + advice())
	}
	return Image{}, ErrNoImage
}

func advice() string {
	switch runtime.GOOS {
	case "darwin":
		return "brew install pngpaste"
	case "windows":
		return "powershell is required"
	default:
		return "install wl-clipboard (Wayland) or xclip (X11)"
	}
}

func run(g grabber) ([]byte, error) {
	ctx, cancel := context.WithTimeout(context.Background(), grabTimeout)
	defer cancel()
	cmd := exec.CommandContext(ctx, g.cmd, g.args...)
	cmd.Stderr = nil
	cmd.Stdin = nil
	// The helpers inherit a terminal they must not write to: this process owns
	// the screen, and a warning printed into the alternate buffer stays there.
	cmd.Env = append(os.Environ(), "NO_COLOR=1")
	return cmd.Output()
}