fckeuspy-go/portal_capture.go

package main

import (
	"errors"
	"fmt"
	"image"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"time"

	"github.com/godbus/dbus/v5"
)

const (
	portalBusName       = "org.freedesktop.portal.Desktop"
	portalObjectPath    = dbus.ObjectPath("/org/freedesktop/portal/desktop")
	screenCastInterface = "org.freedesktop.portal.ScreenCast"
	requestInterface    = "org.freedesktop.portal.Request"
	sessionInterface    = "org.freedesktop.portal.Session"
)

type portalStream struct {
	NodeID     uint32
	Properties map[string]dbus.Variant
}

func openPortalScreenCastSession() (*screenCaptureSession, error) {
	if _, err := exec.LookPath("gst-launch-1.0"); err != nil {
		return nil, errors.New("chybí gst-launch-1.0 pro přenos GNOME ScreenCast")
	}

	conn, err := dbus.ConnectSessionBus()
	if err != nil {
		return nil, fmt.Errorf("připojení k desktop portálu: %w", err)
	}
	closeConnection := true
	defer func() {
		if closeConnection {
			conn.Close()
		}
	}()

	portal := conn.Object(portalBusName, portalObjectPath)
	token := fmt.Sprintf("fckeuspy%d", time.Now().UnixNano())
	createResult, err := portalRequest(conn, portal, screenCastInterface+".CreateSession", map[string]dbus.Variant{
		"handle_token":         dbus.MakeVariant(token + "create"),
		"session_handle_token": dbus.MakeVariant(token + "session"),
	})
	if err != nil {
		return nil, fmt.Errorf("vytvoření ScreenCast relace: %w", err)
	}
	sessionPath, err := portalObjectPathResult(createResult, "session_handle")
	if err != nil {
		return nil, err
	}
	sessionOpen := true
	defer func() {
		if sessionOpen {
			_ = portal.Call(sessionInterface+".Close", 0, sessionPath).Err
		}
	}()

	if _, err := portalRequest(conn, portal, screenCastInterface+".SelectSources", sessionPath, map[string]dbus.Variant{
		"handle_token": dbus.MakeVariant(token + "select"),
		"types":        dbus.MakeVariant(uint32(3)), // monitor or window
		"multiple":     dbus.MakeVariant(false),
		"cursor_mode":  dbus.MakeVariant(uint32(1)), // hidden
	}); err != nil {
		return nil, fmt.Errorf("výběr okna nebo monitoru: %w", err)
	}

	startResult, err := portalRequest(conn, portal, screenCastInterface+".Start", sessionPath, "", map[string]dbus.Variant{
		"handle_token": dbus.MakeVariant(token + "start"),
	})
	if err != nil {
		return nil, fmt.Errorf("spuštění ScreenCast relace: %w", err)
	}
	streams, err := portalStreams(startResult)
	if err != nil {
		return nil, err
	}
	if len(streams) != 1 {
		return nil, fmt.Errorf("ScreenCast vrátil %d streamů, očekáván je jeden", len(streams))
	}

	var remote dbus.UnixFD
	if err := portal.Call(screenCastInterface+".OpenPipeWireRemote", 0, sessionPath, map[string]dbus.Variant{}).Store(&remote); err != nil {
		return nil, fmt.Errorf("otevření PipeWire streamu: %w", err)
	}
	frames, err := openPipeWireFrameStream(remote, streams[0].NodeID)
	if err != nil {
		return nil, err
	}
	closeConnection = false
	sessionOpen = false
	return &screenCaptureSession{
		capture: frames.Capture,
		close: func() {
			frames.Close()
			_ = portal.Call(sessionInterface+".Close", 0, sessionPath).Err
			conn.Close()
		},
	}, nil
}

func portalRequest(conn *dbus.Conn, portal dbus.BusObject, method string, args ...any) (map[string]dbus.Variant, error) {
	signals := make(chan *dbus.Signal, 1)
	conn.Signal(signals)
	defer conn.RemoveSignal(signals)
	if err := conn.AddMatchSignal(
		dbus.WithMatchInterface(requestInterface),
		dbus.WithMatchMember("Response"),
	); err != nil {
		return nil, err
	}
	defer conn.RemoveMatchSignal(
		dbus.WithMatchInterface(requestInterface),
		dbus.WithMatchMember("Response"),
	)

	var requestPath dbus.ObjectPath
	if err := portal.Call(method, 0, args...).Store(&requestPath); err != nil {
		return nil, err
	}

	timeout := time.NewTimer(2 * time.Minute)
	defer timeout.Stop()
	select {
	case <-timeout.C:
		return nil, errors.New("desktop portál neodpověděl včas")
	default:
	}

	for {
		select {
		case signal := <-signals:
			if signal.Path != requestPath {
				continue
			}
			if len(signal.Body) != 2 {
				return nil, errors.New("neplatná odpověď desktop portálu")
			}
			response, ok := signal.Body[0].(uint32)
			if !ok {
				return nil, errors.New("neplatný stav odpovědi desktop portálu")
			}
			if response != 0 {
				return nil, errors.New("výběr byl zrušen")
			}
			result, ok := signal.Body[1].(map[string]dbus.Variant)
			if !ok {
				return nil, errors.New("neplatná data odpovědi desktop portálu")
			}
			return result, nil
		case <-timeout.C:
			return nil, errors.New("desktop portál neodpověděl včas")
		}
	}
}

func portalObjectPathResult(result map[string]dbus.Variant, key string) (dbus.ObjectPath, error) {
	value, ok := result[key]
	if !ok {
		return "", fmt.Errorf("desktop portál nevrátil %s", key)
	}
	var path dbus.ObjectPath
	if err := value.Store(&path); err != nil {
		return "", fmt.Errorf("neplatné %s: %w", key, err)
	}
	return path, nil
}

func portalStreams(result map[string]dbus.Variant) ([]portalStream, error) {
	value, ok := result["streams"]
	if !ok {
		return nil, errors.New("desktop portál nevrátil žádný stream")
	}
	var streams []portalStream
	if err := value.Store(&streams); err != nil {
		return nil, fmt.Errorf("neplatná data streamu desktop portálu: %w", err)
	}
	return streams, nil
}

type pipeWireFrameStream struct {
	dir     string
	command *exec.Cmd
}

func openPipeWireFrameStream(remote dbus.UnixFD, nodeID uint32) (*pipeWireFrameStream, error) {
	dir, err := os.MkdirTemp("", "fckeuspy-pipewire-*")
	if err != nil {
		return nil, err
	}
	remoteFile := os.NewFile(uintptr(remote), "pipewire-remote")
	command := exec.Command(
		"gst-launch-1.0", "-q",
		"pipewiresrc", "fd=3", fmt.Sprintf("path=%d", nodeID),
		"!", "videoconvert",
		"!", "videorate",
		"!", "video/x-raw,framerate=5/1",
		"!", "pngenc", "compression-level=1",
		"!", "multifilesink", "location="+filepath.Join(dir, "frame-%05d.png"), "max-files=2",
	)
	command.ExtraFiles = []*os.File{remoteFile}
	if err := command.Start(); err != nil {
		remoteFile.Close()
		os.Remove(dir)
		return nil, fmt.Errorf("spuštění PipeWire čtečky: %w", err)
	}
	remoteFile.Close()
	return &pipeWireFrameStream{dir: dir, command: command}, nil
}

func (s *pipeWireFrameStream) Capture() (image.Image, error) {
	entries, err := os.ReadDir(s.dir)
	if err != nil {
		return nil, err
	}
	if len(entries) == 0 {
		return nil, errors.New("PipeWire zatím neposkytl žádný snímek")
	}
	sort.Slice(entries, func(i, j int) bool {
		left, leftErr := entries[i].Info()
		right, rightErr := entries[j].Info()
		if leftErr != nil || rightErr != nil {
			return entries[i].Name() > entries[j].Name()
		}
		return left.ModTime().After(right.ModTime())
	})
	return decodeImageFile(filepath.Join(s.dir, entries[0].Name()))
}

func (s *pipeWireFrameStream) Close() {
	if s.command.Process != nil {
		_ = s.command.Process.Kill()
	}
	_, _ = s.command.Process.Wait()
	entries, err := os.ReadDir(s.dir)
	if err == nil {
		for _, entry := range entries {
			_ = os.Remove(filepath.Join(s.dir, entry.Name()))
		}
	}
	_ = os.Remove(s.dir)
}