327 lines
11 KiB
Go
327 lines
11 KiB
Go
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)
|
|
}
|