mirror of
https://github.com/AvengeMedia/DankMaterialShell.git
synced 2025-12-05 21:15:38 -05:00
378 lines
9.9 KiB
Go
378 lines
9.9 KiB
Go
package main
|
|
|
|
import (
|
|
"bytes"
|
|
"fmt"
|
|
"os"
|
|
"os/exec"
|
|
"path/filepath"
|
|
"strings"
|
|
|
|
"github.com/AvengeMedia/DankMaterialShell/core/internal/screenshot"
|
|
"github.com/spf13/cobra"
|
|
)
|
|
|
|
var (
|
|
ssOutputName string
|
|
ssIncludeCursor bool
|
|
ssFormat string
|
|
ssQuality int
|
|
ssOutputDir string
|
|
ssFilename string
|
|
ssNoClipboard bool
|
|
ssNoFile bool
|
|
ssNoNotify bool
|
|
ssStdout bool
|
|
)
|
|
|
|
var screenshotCmd = &cobra.Command{
|
|
Use: "screenshot",
|
|
Short: "Capture screenshots",
|
|
Long: `Capture screenshots from Wayland displays.
|
|
|
|
Modes:
|
|
region - Select a region interactively (default)
|
|
full - Capture the focused output
|
|
all - Capture all outputs combined
|
|
output - Capture a specific output by name
|
|
window - Capture the focused window (Hyprland only)
|
|
last - Capture the last selected region
|
|
|
|
Output format (--format):
|
|
png - PNG format (default)
|
|
jpg/jpeg - JPEG format
|
|
ppm - PPM format
|
|
|
|
Examples:
|
|
dms screenshot # Region select, save file + clipboard
|
|
dms screenshot full # Full screen of focused output
|
|
dms screenshot all # All screens combined
|
|
dms screenshot output -o DP-1 # Specific output
|
|
dms screenshot window # Focused window (Hyprland)
|
|
dms screenshot last # Last region (pre-selected)
|
|
dms screenshot --no-clipboard # Save file only
|
|
dms screenshot --no-file # Clipboard only
|
|
dms screenshot --cursor # Include cursor
|
|
dms screenshot -f jpg -q 85 # JPEG with quality 85`,
|
|
}
|
|
|
|
var ssRegionCmd = &cobra.Command{
|
|
Use: "region",
|
|
Short: "Select a region interactively",
|
|
Run: runScreenshotRegion,
|
|
}
|
|
|
|
var ssFullCmd = &cobra.Command{
|
|
Use: "full",
|
|
Short: "Capture the focused output",
|
|
Run: runScreenshotFull,
|
|
}
|
|
|
|
var ssAllCmd = &cobra.Command{
|
|
Use: "all",
|
|
Short: "Capture all outputs combined",
|
|
Run: runScreenshotAll,
|
|
}
|
|
|
|
var ssOutputCmd = &cobra.Command{
|
|
Use: "output",
|
|
Short: "Capture a specific output",
|
|
Run: runScreenshotOutput,
|
|
}
|
|
|
|
var ssLastCmd = &cobra.Command{
|
|
Use: "last",
|
|
Short: "Capture the last selected region",
|
|
Long: `Capture the previously selected region without interactive selection.
|
|
If no previous region exists, falls back to interactive selection.`,
|
|
Run: runScreenshotLast,
|
|
}
|
|
|
|
var ssWindowCmd = &cobra.Command{
|
|
Use: "window",
|
|
Short: "Capture the focused window",
|
|
Long: `Capture the currently focused window.
|
|
Currently only supported on Hyprland.`,
|
|
Run: runScreenshotWindow,
|
|
}
|
|
|
|
var ssListCmd = &cobra.Command{
|
|
Use: "list",
|
|
Short: "List available outputs",
|
|
Run: runScreenshotList,
|
|
}
|
|
|
|
var notifyActionCmd = &cobra.Command{
|
|
Use: "notify-action",
|
|
Hidden: true,
|
|
Run: func(cmd *cobra.Command, args []string) {
|
|
screenshot.RunNotifyActionListener(args)
|
|
},
|
|
}
|
|
|
|
func init() {
|
|
screenshotCmd.PersistentFlags().StringVarP(&ssOutputName, "output", "o", "", "Output name for 'output' mode")
|
|
screenshotCmd.PersistentFlags().BoolVar(&ssIncludeCursor, "cursor", false, "Include cursor in screenshot")
|
|
screenshotCmd.PersistentFlags().StringVarP(&ssFormat, "format", "f", "png", "Output format (png, jpg, ppm)")
|
|
screenshotCmd.PersistentFlags().IntVarP(&ssQuality, "quality", "q", 90, "JPEG quality (1-100)")
|
|
screenshotCmd.PersistentFlags().StringVarP(&ssOutputDir, "dir", "d", "", "Output directory")
|
|
screenshotCmd.PersistentFlags().StringVar(&ssFilename, "filename", "", "Output filename (auto-generated if empty)")
|
|
screenshotCmd.PersistentFlags().BoolVar(&ssNoClipboard, "no-clipboard", false, "Don't copy to clipboard")
|
|
screenshotCmd.PersistentFlags().BoolVar(&ssNoFile, "no-file", false, "Don't save to file")
|
|
screenshotCmd.PersistentFlags().BoolVar(&ssNoNotify, "no-notify", false, "Don't show notification")
|
|
screenshotCmd.PersistentFlags().BoolVar(&ssStdout, "stdout", false, "Output image to stdout (for piping to swappy, etc.)")
|
|
|
|
screenshotCmd.AddCommand(ssRegionCmd)
|
|
screenshotCmd.AddCommand(ssFullCmd)
|
|
screenshotCmd.AddCommand(ssAllCmd)
|
|
screenshotCmd.AddCommand(ssOutputCmd)
|
|
screenshotCmd.AddCommand(ssLastCmd)
|
|
screenshotCmd.AddCommand(ssWindowCmd)
|
|
screenshotCmd.AddCommand(ssListCmd)
|
|
|
|
screenshotCmd.Run = runScreenshotRegion
|
|
}
|
|
|
|
func getScreenshotConfig(mode screenshot.Mode) screenshot.Config {
|
|
config := screenshot.DefaultConfig()
|
|
config.Mode = mode
|
|
config.OutputName = ssOutputName
|
|
config.IncludeCursor = ssIncludeCursor
|
|
config.Clipboard = !ssNoClipboard
|
|
config.SaveFile = !ssNoFile
|
|
config.Notify = !ssNoNotify
|
|
config.Stdout = ssStdout
|
|
|
|
if ssOutputDir != "" {
|
|
config.OutputDir = ssOutputDir
|
|
}
|
|
if ssFilename != "" {
|
|
config.Filename = ssFilename
|
|
}
|
|
|
|
switch strings.ToLower(ssFormat) {
|
|
case "jpg", "jpeg":
|
|
config.Format = screenshot.FormatJPEG
|
|
case "ppm":
|
|
config.Format = screenshot.FormatPPM
|
|
default:
|
|
config.Format = screenshot.FormatPNG
|
|
}
|
|
|
|
if ssQuality < 1 {
|
|
ssQuality = 1
|
|
}
|
|
if ssQuality > 100 {
|
|
ssQuality = 100
|
|
}
|
|
config.Quality = ssQuality
|
|
|
|
return config
|
|
}
|
|
|
|
func runScreenshot(config screenshot.Config) {
|
|
sc := screenshot.New(config)
|
|
result, err := sc.Run()
|
|
if err != nil {
|
|
fmt.Fprintf(os.Stderr, "Error: %v\n", err)
|
|
os.Exit(1)
|
|
}
|
|
|
|
if result == nil {
|
|
os.Exit(0)
|
|
}
|
|
|
|
defer result.Buffer.Close()
|
|
|
|
if result.YInverted {
|
|
result.Buffer.FlipVertical()
|
|
}
|
|
|
|
if config.Stdout {
|
|
if err := writeImageToStdout(result.Buffer, config.Format, config.Quality, result.Format); err != nil {
|
|
fmt.Fprintf(os.Stderr, "Error writing to stdout: %v\n", err)
|
|
os.Exit(1)
|
|
}
|
|
return
|
|
}
|
|
|
|
var filePath string
|
|
|
|
if config.SaveFile {
|
|
outputDir := config.OutputDir
|
|
if outputDir == "" {
|
|
outputDir = screenshot.GetOutputDir()
|
|
}
|
|
|
|
filename := config.Filename
|
|
if filename == "" {
|
|
filename = screenshot.GenerateFilename(config.Format)
|
|
}
|
|
|
|
filePath = filepath.Join(outputDir, filename)
|
|
if err := screenshot.WriteToFileWithFormat(result.Buffer, filePath, config.Format, config.Quality, result.Format); err != nil {
|
|
fmt.Fprintf(os.Stderr, "Error writing file: %v\n", err)
|
|
os.Exit(1)
|
|
}
|
|
fmt.Println(filePath)
|
|
}
|
|
|
|
if config.Clipboard {
|
|
if err := copyImageToClipboard(result.Buffer, config.Format, config.Quality, result.Format); err != nil {
|
|
fmt.Fprintf(os.Stderr, "Error copying to clipboard: %v\n", err)
|
|
os.Exit(1)
|
|
}
|
|
if !config.SaveFile {
|
|
fmt.Println("Copied to clipboard")
|
|
}
|
|
}
|
|
|
|
if config.Notify {
|
|
thumbData, thumbW, thumbH := bufferToRGBThumbnail(result.Buffer, 256, result.Format)
|
|
screenshot.SendNotification(screenshot.NotifyResult{
|
|
FilePath: filePath,
|
|
Clipboard: config.Clipboard,
|
|
ImageData: thumbData,
|
|
Width: thumbW,
|
|
Height: thumbH,
|
|
})
|
|
}
|
|
}
|
|
|
|
func copyImageToClipboard(buf *screenshot.ShmBuffer, format screenshot.Format, quality int, pixelFormat uint32) error {
|
|
var mimeType string
|
|
var data bytes.Buffer
|
|
|
|
img := screenshot.BufferToImageWithFormat(buf, pixelFormat)
|
|
|
|
switch format {
|
|
case screenshot.FormatJPEG:
|
|
mimeType = "image/jpeg"
|
|
if err := screenshot.EncodeJPEG(&data, img, quality); err != nil {
|
|
return err
|
|
}
|
|
default:
|
|
mimeType = "image/png"
|
|
if err := screenshot.EncodePNG(&data, img); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
cmd := exec.Command("wl-copy", "--type", mimeType)
|
|
cmd.Stdin = &data
|
|
return cmd.Run()
|
|
}
|
|
|
|
func writeImageToStdout(buf *screenshot.ShmBuffer, format screenshot.Format, quality int, pixelFormat uint32) error {
|
|
img := screenshot.BufferToImageWithFormat(buf, pixelFormat)
|
|
|
|
switch format {
|
|
case screenshot.FormatJPEG:
|
|
return screenshot.EncodeJPEG(os.Stdout, img, quality)
|
|
default:
|
|
return screenshot.EncodePNG(os.Stdout, img)
|
|
}
|
|
}
|
|
|
|
func bufferToRGBThumbnail(buf *screenshot.ShmBuffer, maxSize int, pixelFormat uint32) ([]byte, int, int) {
|
|
srcW, srcH := buf.Width, buf.Height
|
|
scale := 1.0
|
|
if srcW > maxSize || srcH > maxSize {
|
|
if srcW > srcH {
|
|
scale = float64(maxSize) / float64(srcW)
|
|
} else {
|
|
scale = float64(maxSize) / float64(srcH)
|
|
}
|
|
}
|
|
|
|
dstW := int(float64(srcW) * scale)
|
|
dstH := int(float64(srcH) * scale)
|
|
if dstW < 1 {
|
|
dstW = 1
|
|
}
|
|
if dstH < 1 {
|
|
dstH = 1
|
|
}
|
|
|
|
data := buf.Data()
|
|
rgb := make([]byte, dstW*dstH*3)
|
|
swapRB := pixelFormat == uint32(screenshot.FormatARGB8888) || pixelFormat == uint32(screenshot.FormatXRGB8888) || pixelFormat == 0
|
|
|
|
for y := 0; y < dstH; y++ {
|
|
srcY := int(float64(y) / scale)
|
|
if srcY >= srcH {
|
|
srcY = srcH - 1
|
|
}
|
|
for x := 0; x < dstW; x++ {
|
|
srcX := int(float64(x) / scale)
|
|
if srcX >= srcW {
|
|
srcX = srcW - 1
|
|
}
|
|
si := srcY*buf.Stride + srcX*4
|
|
di := (y*dstW + x) * 3
|
|
if si+2 < len(data) {
|
|
if swapRB {
|
|
rgb[di+0] = data[si+2]
|
|
rgb[di+1] = data[si+1]
|
|
rgb[di+2] = data[si+0]
|
|
} else {
|
|
rgb[di+0] = data[si+0]
|
|
rgb[di+1] = data[si+1]
|
|
rgb[di+2] = data[si+2]
|
|
}
|
|
}
|
|
}
|
|
}
|
|
return rgb, dstW, dstH
|
|
}
|
|
|
|
func runScreenshotRegion(cmd *cobra.Command, args []string) {
|
|
config := getScreenshotConfig(screenshot.ModeRegion)
|
|
runScreenshot(config)
|
|
}
|
|
|
|
func runScreenshotFull(cmd *cobra.Command, args []string) {
|
|
config := getScreenshotConfig(screenshot.ModeFullScreen)
|
|
runScreenshot(config)
|
|
}
|
|
|
|
func runScreenshotAll(cmd *cobra.Command, args []string) {
|
|
config := getScreenshotConfig(screenshot.ModeAllScreens)
|
|
runScreenshot(config)
|
|
}
|
|
|
|
func runScreenshotOutput(cmd *cobra.Command, args []string) {
|
|
if ssOutputName == "" && len(args) > 0 {
|
|
ssOutputName = args[0]
|
|
}
|
|
if ssOutputName == "" {
|
|
fmt.Fprintln(os.Stderr, "Error: output name required (use -o or provide as argument)")
|
|
os.Exit(1)
|
|
}
|
|
config := getScreenshotConfig(screenshot.ModeOutput)
|
|
runScreenshot(config)
|
|
}
|
|
|
|
func runScreenshotLast(cmd *cobra.Command, args []string) {
|
|
config := getScreenshotConfig(screenshot.ModeLastRegion)
|
|
runScreenshot(config)
|
|
}
|
|
|
|
func runScreenshotWindow(cmd *cobra.Command, args []string) {
|
|
config := getScreenshotConfig(screenshot.ModeWindow)
|
|
runScreenshot(config)
|
|
}
|
|
|
|
func runScreenshotList(cmd *cobra.Command, args []string) {
|
|
outputs, err := screenshot.ListOutputs()
|
|
if err != nil {
|
|
fmt.Fprintf(os.Stderr, "Error: %v\n", err)
|
|
os.Exit(1)
|
|
}
|
|
|
|
for _, o := range outputs {
|
|
fmt.Printf("%s: %dx%d+%d+%d (scale: %d)\n",
|
|
o.Name, o.Width, o.Height, o.X, o.Y, o.Scale)
|
|
}
|
|
}
|