wayland-screenshot/capture.go
2023-02-06 18:14:14 +03:00

175 lines
3.2 KiB
Go
Executable File

package main
import (
"bytes"
"errors"
"mime"
"net/http"
"os"
"os/exec"
"path/filepath"
"time"
)
const FormatFileDefault = "2006-01-02_15-04-05"
type Capture int
var ErrUnknownCaptureType = errors.New("unknown capture type")
const (
Area Capture = iota
All
)
// ParseCapture parses a string to a Capture type
func ParseCapture(s string) (Capture, error) {
switch s {
case "area":
return Area, nil
case "all":
return All, nil
}
return 0, ErrUnknownCaptureType
}
func (c Capture) String() string {
switch c {
case Area:
return "area"
case All:
return "all"
}
return ""
}
// CaptureArea captures a region
func CaptureArea(r Region) (Captured, error) {
output, err := exec.Command("grim", "-g", r.String(), "-").Output() //nolint:gosec // dn
return Captured{output}, err
}
// CaptureAll captures the whole screen
func CaptureAll() (Captured, error) {
output, err := exec.Command("grim", "-").Output()
return Captured{output}, err
}
// Captured is a captured image
type Captured struct {
Raw []byte
}
// SaveFile saves the captured image to a file
func (c *Captured) SaveFile(path string) error {
return saveFileToDir(c.Raw, path)
}
// SaveClipboard saves the captured image to the clipboard
func (c *Captured) SaveClipboard() error {
cmd := exec.Command("wl-copy")
cmd.Stderr = os.Stderr
cmd.Stdin = bytes.NewReader(c.Raw)
return cmd.Run()
}
// Edit edits the captured image with swappy
func (c *Captured) Edit() (Captured, error) {
cmd := exec.Command("swappy", "-f", "-", "-o", "-")
cmd.Stderr = os.Stderr
cmd.Stdin = bytes.NewReader(c.Raw)
var out bytes.Buffer
cmd.Stdout = &out
err := cmd.Run()
if err != nil {
return Captured{}, err
}
return Captured{out.Bytes()}, nil
}
// generateNameDefaultFormat generates a name for a file with the default format
func generateNameDefaultFormat(raw []byte) (string, error) {
ext, err := getExt(raw)
if err != nil {
return "", err
}
name := time.Now().Format(FormatFileDefault) + ext
return name, nil
}
// getExt returns the extension of a mime type
func getExt(raw []byte) (string, error) {
byType, mimeTypeErr := mime.ExtensionsByType(http.DetectContentType(raw))
if mimeTypeErr != nil {
return "", mimeTypeErr
}
if len(byType) == 0 {
return "", errors.New("unknown file type")
}
return byType[0], nil
}
// saveFileToDir saves a file to a directory
func saveFileToDir(raw []byte, path string) error {
if err := ensureDir(path); err != nil {
return err
}
p, err := isPath(path)
if err != nil {
return err
}
if p {
name, genErr := generateNameDefaultFormat(raw)
if genErr != nil {
return err
}
path = path + "/" + name
}
file, createErr := os.Create(path)
if createErr != nil {
return createErr
}
_, err = file.Write(raw)
return err
}
// ensureDir create recursively a directory
func ensureDir(fileName string) error {
dirName := filepath.Dir(fileName)
if _, serr := os.Stat(dirName); serr != nil {
return os.MkdirAll(dirName, os.ModePerm)
}
return nil
}
// isPath checks if a path is a directory or a file
func isPath(path string) (bool, error) {
stat, err := os.Stat(path)
if err != nil {
// if no such file or directory, is ok
if os.IsNotExist(err) {
return false, nil
}
return false, err
}
return stat.IsDir(), nil
}