package agent import ( "context" "encoding/json" "errors" "io" "mime" "net/http" "os" "path/filepath" "runtime" "strconv" "strings" "time" "tea.chunkbyte.com/kato/go-worm/lib/command" "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/models" "tea.chunkbyte.com/kato/go-worm/lib/screenshot" "tea.chunkbyte.com/kato/go-worm/lib/startup" ) func (a *Agent) handleHealth(w http.ResponseWriter, r *http.Request) { if r.Method != http.MethodGet { helpers.WriteError(w, http.StatusMethodNotAllowed, "method not allowed") return } helpers.WriteJSON(w, http.StatusOK, map[string]string{"status": "ok"}) } func (a *Agent) handleOpenAPI(w http.ResponseWriter, r *http.Request) { if r.Method != http.MethodGet { helpers.WriteError(w, http.StatusMethodNotAllowed, "method not allowed") return } helpers.WriteJSON(w, http.StatusOK, map[string]any{ "openapi": "3.0.3", "info": map[string]string{"title": "Local Management Agent", "version": config.Version}, "paths": map[string]any{ "/api/v1/status": map[string]any{"get": map[string]string{"summary": "Agent status"}}, "/api/v1/files": map[string]any{"get": map[string]string{"summary": "List files"}}, "/api/v1/download": map[string]any{"get": map[string]string{"summary": "Download file"}}, "/api/v1/screenshot": map[string]any{"get": map[string]string{"summary": "Capture desktop"}}, "/api/v1/exec": map[string]any{"post": map[string]string{"summary": "Run a command"}}, "/api/v1/startup": map[string]any{"post": map[string]string{"summary": "Add to Windows startup"}, "delete": map[string]string{"summary": "Remove from Windows startup"}}, "/api/v1/input/click": map[string]any{"post": map[string]string{"summary": "Click the desktop"}}, "/api/v1/input/text": map[string]any{"post": map[string]string{"summary": "Type text into the focused field"}}, }, }) } func (a *Agent) handleStatus(w http.ResponseWriter, r *http.Request) { if r.Method != http.MethodGet { helpers.WriteError(w, http.StatusMethodNotAllowed, "method not allowed") return } host, _ := os.Hostname() helpers.WriteJSON(w, http.StatusOK, map[string]any{ "os": "windows", "architecture": runtime.GOARCH, "user": helpers.Username(), "hostname": host, "uptime_seconds": int64(time.Since(a.startedAt).Seconds()), "local_ips": helpers.LocalIPs(), "agent_version": config.Version, "listen_address": a.addr, "startup_enabled": startup.Enabled(), }) } func (a *Agent) handleStartup(w http.ResponseWriter, r *http.Request) { switch r.Method { case http.MethodPost: if err := startup.Enable(); err != nil { helpers.Log.Printf("startup enable: %v", err) helpers.WriteError(w, http.StatusInternalServerError, "could not add to startup") return } case http.MethodDelete: if err := startup.Disable(); err != nil { helpers.Log.Printf("startup disable: %v", err) helpers.WriteError(w, http.StatusInternalServerError, "could not remove from startup") return } default: helpers.WriteError(w, http.StatusMethodNotAllowed, "method not allowed") return } helpers.WriteJSON(w, http.StatusOK, map[string]any{"startup_enabled": startup.Enabled()}) } func (a *Agent) handleFiles(w http.ResponseWriter, r *http.Request) { if r.Method != http.MethodGet { helpers.WriteError(w, http.StatusMethodNotAllowed, "method not allowed") return } depth, err := files.ParseDepth(r.URL.Query().Get("depth")) if err != nil { helpers.WriteError(w, http.StatusBadRequest, err.Error()) return } path := r.URL.Query().Get("path") if path == "" { path, _ = os.UserHomeDir() } dir, err := files.AllowedPath(a.root, path, true) if err != nil { files.WritePathError(w, err) return } info, err := os.Stat(dir) if err != nil { files.WritePathError(w, err) return } if !info.IsDir() { helpers.WriteError(w, http.StatusBadRequest, "path is not a directory") return } entries, err := files.ListDirectory(dir, depth) if err != nil { files.WritePathError(w, err) return } helpers.WriteJSON(w, http.StatusOK, map[string]any{"path": dir, "depth": depth, "entries": entries}) } func (a *Agent) handleDownload(w http.ResponseWriter, r *http.Request) { if r.Method != http.MethodGet { helpers.WriteError(w, http.StatusMethodNotAllowed, "method not allowed") return } path := r.URL.Query().Get("path") if path == "" { helpers.WriteError(w, http.StatusBadRequest, "path is required") return } file, err := files.AllowedPath(a.root, path, true) if err != nil { files.WritePathError(w, err) return } f, err := os.Open(file) if err != nil { files.WritePathError(w, err) return } defer f.Close() info, err := f.Stat() if err != nil { files.WritePathError(w, err) return } if info.IsDir() { helpers.WriteError(w, http.StatusBadRequest, "path is a directory") return } name := filepath.Base(file) contentType := mime.TypeByExtension(filepath.Ext(name)) if contentType == "" { var sample [512]byte n, _ := f.Read(sample[:]) contentType = http.DetectContentType(sample[:n]) _, _ = f.Seek(0, io.SeekStart) } w.Header().Set("Content-Type", contentType) w.Header().Set("Content-Disposition", `attachment; filename="`+strings.ReplaceAll(name, `"`, "'")+`"`) w.Header().Set("Accept-Ranges", "bytes") http.ServeContent(w, r, name, info.ModTime(), f) } func (a *Agent) handleScreenshot(w http.ResponseWriter, r *http.Request) { if r.Method != http.MethodGet { helpers.WriteError(w, http.StatusMethodNotAllowed, "method not allowed") return } format := strings.ToLower(strings.TrimSpace(r.URL.Query().Get("format"))) if format == "" { format = "png" } if format != "png" && format != "jpeg" { helpers.WriteError(w, http.StatusBadRequest, "format must be png or jpeg") return } quality := 80 if raw := r.URL.Query().Get("quality"); raw != "" { var err error quality, err = strconv.Atoi(raw) if err != nil || quality < 1 || quality > config.MaxImageQuality { helpers.WriteError(w, http.StatusBadRequest, "quality must be between 1 and 100") return } } monitor := 0 if raw := r.URL.Query().Get("monitor"); raw != "" { var err error monitor, err = strconv.Atoi(raw) if err != nil || monitor < 0 { helpers.WriteError(w, http.StatusBadRequest, "monitor must be 0 or greater") return } } frame, err := screenshot.CaptureMonitor(monitor, format, quality) if err != nil { if errors.Is(err, screenshot.ErrMonitorNotFound) { helpers.WriteError(w, http.StatusBadRequest, err.Error()) return } helpers.Log.Printf("screenshot: %v", err) helpers.WriteError(w, http.StatusServiceUnavailable, "no interactive desktop is available") return } w.Header().Set("Content-Type", frame.ContentType) w.Header().Set("Content-Length", strconv.Itoa(len(frame.Data))) w.Header().Set("X-Monitor-Left", strconv.Itoa(frame.Left)) w.Header().Set("X-Monitor-Top", strconv.Itoa(frame.Top)) w.Header().Set("X-Monitor-Width", strconv.Itoa(frame.Width)) w.Header().Set("X-Monitor-Height", strconv.Itoa(frame.Height)) _, _ = w.Write(frame.Data) } func (a *Agent) handleClick(w http.ResponseWriter, r *http.Request) { if r.Method != http.MethodPost { helpers.WriteError(w, http.StatusMethodNotAllowed, "method not allowed") return } r.Body = http.MaxBytesReader(w, r.Body, config.RequestBodyMax) defer r.Body.Close() var request models.ClickRequest decoder := json.NewDecoder(r.Body) decoder.DisallowUnknownFields() if err := decoder.Decode(&request); err != nil { helpers.WriteError(w, http.StatusBadRequest, "body must contain click coordinates") return } left, top, width, height, err := screenshot.MonitorBounds(request.Monitor) if err != nil { if errors.Is(err, screenshot.ErrMonitorNotFound) { helpers.WriteError(w, http.StatusBadRequest, err.Error()) return } helpers.WriteError(w, http.StatusBadRequest, "monitor not found") return } if request.X < left || request.Y < top || request.X >= left+width || request.Y >= top+height { helpers.WriteError(w, http.StatusBadRequest, "click is outside the selected monitor") return } if err := input.Click(request.X, request.Y, request.Button); err != nil { if errors.Is(err, input.ErrBadButton) { helpers.WriteError(w, http.StatusBadRequest, err.Error()) return } helpers.Log.Printf("click: %v", err) helpers.WriteError(w, http.StatusInternalServerError, "could not click") return } helpers.WriteJSON(w, http.StatusOK, map[string]any{"ok": true, "x": request.X, "y": request.Y}) } func (a *Agent) handleText(w http.ResponseWriter, r *http.Request) { if r.Method != http.MethodPost { helpers.WriteError(w, http.StatusMethodNotAllowed, "method not allowed") return } r.Body = http.MaxBytesReader(w, r.Body, config.RequestBodyMax) defer r.Body.Close() var request models.TextRequest decoder := json.NewDecoder(r.Body) decoder.DisallowUnknownFields() if err := decoder.Decode(&request); err != nil { helpers.WriteError(w, http.StatusBadRequest, "body must contain text") return } if request.Text == "" { helpers.WriteError(w, http.StatusBadRequest, input.ErrEmptyText.Error()) return } if len(request.Text) > config.MaxInputText { helpers.WriteError(w, http.StatusBadRequest, "text is too long") return } if err := input.TypeText(request.Text); err != nil { helpers.Log.Printf("text: %v", err) helpers.WriteError(w, http.StatusInternalServerError, "could not type text") return } helpers.WriteJSON(w, http.StatusOK, map[string]any{"ok": true, "length": len(request.Text)}) } func (a *Agent) handleExec(w http.ResponseWriter, r *http.Request) { if r.Method != http.MethodPost { helpers.WriteError(w, http.StatusMethodNotAllowed, "method not allowed") return } r.Body = http.MaxBytesReader(w, r.Body, config.RequestBodyMax) defer r.Body.Close() var request models.ExecRequest decoder := json.NewDecoder(r.Body) decoder.DisallowUnknownFields() if err := decoder.Decode(&request); err != nil { helpers.WriteError(w, http.StatusBadRequest, "body must contain a command") return } cmdline := strings.TrimSpace(request.Command) if cmdline == "" { helpers.WriteError(w, http.StatusBadRequest, command.ErrEmptyCommand.Error()) return } timeout, err := command.ResolveTimeout(request.TimeoutSec) if err != nil { helpers.WriteError(w, http.StatusBadRequest, err.Error()) return } ctx, cancel := context.WithTimeout(r.Context(), timeout) defer cancel() result, err := command.Run(ctx, cmdline) if errors.Is(err, command.ErrTimeout) { helpers.WriteError(w, http.StatusGatewayTimeout, err.Error()) return } if err != nil { helpers.Log.Printf("exec: %v", err) helpers.WriteError(w, http.StatusInternalServerError, "could not run command") return } helpers.WriteJSON(w, http.StatusOK, result) }