TaskSchedulerService.cs

using System.Security.Principal;
using System.Globalization;

namespace BatteryDown;

internal static class TaskSchedulerService
{
    private const string UnplugTaskName = "Battery Down - Power disconnected";
    private const string ReplugTaskName = "Battery Down - Power connected";
    private const string UnlockTaskName = "Battery Down - Session unlocked";
    private const string LockedSystemTaskName = "Battery Down - Locked power check";
    private const string PendingShutdownTaskName = "Battery Down - Pending shutdown";
    private const int TaskCreateOrUpdate = 6;
    private const int TaskLogonInteractiveToken = 3;
    private const int TaskLogonServiceAccount = 5;
    private const int AccessDeniedHResult = unchecked((int)0x80070005);

    public static bool IsLockedSystemTaskInstalled => TaskExists(LockedSystemTaskName);

    public static bool IsInstalled
    {
        get
        {
            dynamic? service = null;
            dynamic? root = null;
            try
            {
                service = CreateService();
                root = service.GetFolder("\\");
                _ = root.GetTask(UnplugTaskName);
                _ = root.GetTask(ReplugTaskName);
                _ = root.GetTask(UnlockTaskName);
                return true;
            }
            catch
            {
                return false;
            }
            finally
            {
                ReleaseComObject(root);
                ReleaseComObject(service);
            }
        }
    }

    public static void Install(AppConfig config)
    {
        RegistrationService.EnsureRegistered();

        var executablePath = Environment.ProcessPath
            ?? throw new InvalidOperationException("无法确定程序路径。");

        dynamic? service = null;
        dynamic? root = null;
        try
        {
            service = CreateService();
            root = service.GetFolder("\\");

            RegisterEventTask(
                service,
                root,
                UnplugTaskName,
                "拔出电源后，延迟启动 Battery Down 提醒。",
                PowerEventSubscription(acOnline: false),
                config.DelaySeconds == 0 ? null : $"PT{config.DelaySeconds}S",
                executablePath,
                "--prompt");

            RegisterEventTask(
                service,
                root,
                ReplugTaskName,
                "重新接通电源后，取消 Battery Down 提醒。",
                PowerEventSubscription(acOnline: true),
                null,
                executablePath,
                "--cancel");

            RegisterSessionUnlockTask(service, root, executablePath);
        }
        finally
        {
            ReleaseComObject(root);
            ReleaseComObject(service);
        }

        SystemLockTaskService.Configure(config.ShutdownImmediatelyWhenLocked);
    }

    public static void Uninstall()
    {
        SystemLockTaskService.Configure(enabled: false);
        ShutdownPlanService.Cancel();

        dynamic? service = null;
        dynamic? root = null;
        try
        {
            service = CreateService();
            root = service.GetFolder("\\");
            DeleteTaskIfPresent(root, UnplugTaskName);
            DeleteTaskIfPresent(root, ReplugTaskName);
            DeleteTaskIfPresent(root, UnlockTaskName);
            DeleteTaskIfPresent(root, PendingShutdownTaskName);
        }
        finally
        {
            ReleaseComObject(root);
            ReleaseComObject(service);
        }

        ToastService.ClearPrompt();
        RegistrationService.RemoveRegistration();
    }

    public static void InstallLockedSystemTask()
    {
        var executablePath = Environment.ProcessPath
            ?? throw new InvalidOperationException("无法确定程序路径。");

        dynamic? service = null;
        dynamic? root = null;
        dynamic? definition = null;
        dynamic? trigger = null;
        dynamic? action = null;
        try
        {
            service = CreateService();
            root = service.GetFolder("\\");
            definition = service.NewTask(0);

            definition.RegistrationInfo.Author = "Battery Down";
            definition.RegistrationInfo.Description =
                "锁屏状态下检测到拔出电源时立即关闭电脑。";
            definition.Principal.UserId = "SYSTEM";
            definition.Principal.LogonType = TaskLogonServiceAccount;
            definition.Principal.RunLevel = 1;
            ConfigureSettings(definition, "PT2M", startWhenAvailable: false);

            trigger = definition.Triggers.Create(0);
            trigger.Enabled = true;
            trigger.Subscription = PowerEventSubscription(acOnline: false);

            action = definition.Actions.Create(0);
            action.Path = executablePath;
            action.Arguments = $"--locked-power-check \"{AppConfig.FilePath}\"";
            action.WorkingDirectory = Path.GetDirectoryName(executablePath);

            root.RegisterTaskDefinition(
                LockedSystemTaskName,
                definition,
                TaskCreateOrUpdate,
                "SYSTEM",
                null,
                TaskLogonServiceAccount,
                null);
        }
        finally
        {
            ReleaseComObject(action);
            ReleaseComObject(trigger);
            ReleaseComObject(definition);
            ReleaseComObject(root);
            ReleaseComObject(service);
        }
    }

    public static void RemoveLockedSystemTask()
    {
        dynamic? service = null;
        dynamic? root = null;
        try
        {
            service = CreateService();
            root = service.GetFolder("\\");
            DeleteTaskIfPresent(root, LockedSystemTaskName);
        }
        finally
        {
            ReleaseComObject(root);
            ReleaseComObject(service);
        }
    }

    public static void SchedulePendingShutdown(string token, DateTimeOffset deadlineUtc)
    {
        var executablePath = Environment.ProcessPath
            ?? throw new InvalidOperationException("无法确定程序路径。");

        dynamic? service = null;
        dynamic? root = null;
        dynamic? definition = null;
        dynamic? trigger = null;
        dynamic? action = null;
        try
        {
            service = CreateService();
            root = service.GetFolder("\\");
            definition = service.NewTask(0);
            ConfigureDefinition(
                definition,
                "Battery Down 倒计时结束后执行关机。",
                "PT2M",
                startWhenAvailable: true);

            trigger = definition.Triggers.Create(1); // TASK_TRIGGER_TIME
            trigger.Enabled = true;
            trigger.StartBoundary = deadlineUtc.LocalDateTime.ToString(
                "yyyy-MM-dd'T'HH:mm:ss",
                CultureInfo.InvariantCulture);

            action = definition.Actions.Create(0);
            action.Path = executablePath;
            action.Arguments = $"--timeout {token}";
            action.WorkingDirectory = Path.GetDirectoryName(executablePath);

            root.RegisterTaskDefinition(
                PendingShutdownTaskName,
                definition,
                TaskCreateOrUpdate,
                WindowsIdentity.GetCurrent().Name,
                null,
                TaskLogonInteractiveToken,
                null);
        }
        finally
        {
            ReleaseComObject(action);
            ReleaseComObject(trigger);
            ReleaseComObject(definition);
            ReleaseComObject(root);
            ReleaseComObject(service);
        }
    }

    public static void CancelPendingShutdown()
    {
        dynamic? service = null;
        dynamic? root = null;
        try
        {
            service = CreateService();
            root = service.GetFolder("\\");
            DeleteTaskIfPresent(root, PendingShutdownTaskName);
        }
        finally
        {
            ReleaseComObject(root);
            ReleaseComObject(service);
        }
    }

    private static dynamic CreateService()
    {
        var serviceType = Type.GetTypeFromProgID("Schedule.Service")
            ?? throw new PlatformNotSupportedException("当前系统不支持 Windows 任务计划程序。");
        dynamic service = Activator.CreateInstance(serviceType)
            ?? throw new InvalidOperationException("无法连接 Windows 任务计划程序。");
        service.Connect();
        return service;
    }

    private static void RegisterEventTask(
        dynamic service,
        dynamic root,
        string taskName,
        string description,
        string eventSubscription,
        string? delay,
        string executablePath,
        string arguments)
    {
        dynamic? definition = null;
        dynamic? trigger = null;
        dynamic? action = null;
        try
        {
            definition = service.NewTask(0);
            ConfigureDefinition(definition, description, "PT2H", startWhenAvailable: false);

            trigger = definition.Triggers.Create(0); // TASK_TRIGGER_EVENT
            trigger.Enabled = true;
            trigger.Subscription = eventSubscription;
            if (!string.IsNullOrWhiteSpace(delay))
            {
                trigger.Delay = delay;
            }

            action = definition.Actions.Create(0); // TASK_ACTION_EXEC
            action.Path = executablePath;
            action.Arguments = arguments;
            action.WorkingDirectory = Path.GetDirectoryName(executablePath);

            root.RegisterTaskDefinition(
                taskName,
                definition,
                TaskCreateOrUpdate,
                WindowsIdentity.GetCurrent().Name,
                null,
                TaskLogonInteractiveToken,
                null);
        }
        finally
        {
            ReleaseComObject(action);
            ReleaseComObject(trigger);
            ReleaseComObject(definition);
        }
    }

    private static void RegisterSessionUnlockTask(dynamic service, dynamic root, string executablePath)
    {
        dynamic? definition = null;
        dynamic? trigger = null;
        dynamic? action = null;
        try
        {
            definition = service.NewTask(0);
            ConfigureDefinition(
                definition,
                "解锁 Windows 会话后，如果仍在使用电池则显示 Battery Down 提醒。",
                "PT2H",
                startWhenAvailable: false);

            trigger = definition.Triggers.Create(11); // TASK_TRIGGER_SESSION_STATE_CHANGE
            trigger.Enabled = true;
            trigger.StateChange = 8; // TASK_SESSION_UNLOCK
            trigger.UserId = WindowsIdentity.GetCurrent().Name;

            action = definition.Actions.Create(0);
            action.Path = executablePath;
            action.Arguments = "--unlock";
            action.WorkingDirectory = Path.GetDirectoryName(executablePath);

            root.RegisterTaskDefinition(
                UnlockTaskName,
                definition,
                TaskCreateOrUpdate,
                WindowsIdentity.GetCurrent().Name,
                null,
                TaskLogonInteractiveToken,
                null);
        }
        finally
        {
            ReleaseComObject(action);
            ReleaseComObject(trigger);
            ReleaseComObject(definition);
        }
    }

    private static void ConfigureDefinition(
        dynamic definition,
        string description,
        string executionTimeLimit,
        bool startWhenAvailable)
    {
        definition.RegistrationInfo.Author = "Battery Down";
        definition.RegistrationInfo.Description = description;

        definition.Principal.UserId = WindowsIdentity.GetCurrent().Name;
        definition.Principal.LogonType = TaskLogonInteractiveToken;
        definition.Principal.RunLevel = 0;

        ConfigureSettings(definition, executionTimeLimit, startWhenAvailable);
    }

    private static void ConfigureSettings(
        dynamic definition,
        string executionTimeLimit,
        bool startWhenAvailable)
    {

        definition.Settings.Enabled = true;
        definition.Settings.Hidden = false;
        definition.Settings.AllowDemandStart = true;
        definition.Settings.StartWhenAvailable = startWhenAvailable;
        definition.Settings.DisallowStartIfOnBatteries = false;
        definition.Settings.StopIfGoingOnBatteries = false;
        definition.Settings.ExecutionTimeLimit = executionTimeLimit;
        definition.Settings.MultipleInstances = 2;
    }

    private static bool TaskExists(string taskName)
    {
        dynamic? service = null;
        dynamic? root = null;
        try
        {
            service = CreateService();
            root = service.GetFolder("\\");
            _ = root.GetTask(taskName);
            return true;
        }
        catch (Exception exception) when (exception.HResult == AccessDeniedHResult)
        {
            // SYSTEM tasks use a protected ACL. Access denied means the named task exists,
            // but the current non-elevated process cannot read its definition.
            return true;
        }
        catch
        {
            return false;
        }
        finally
        {
            ReleaseComObject(root);
            ReleaseComObject(service);
        }
    }

    private static string PowerEventSubscription(bool acOnline)
    {
        var value = acOnline ? "true" : "false";
        return $"<QueryList><Query Id='0' Path='System'><Select Path='System'>" +
               "*[System[Provider[@Name='Microsoft-Windows-Kernel-Power'] and EventID=105]] " +
               $"and *[EventData[Data[@Name='AcOnline']='{value}']]" +
               "</Select></Query></QueryList>";
    }

    private static void DeleteTaskIfPresent(dynamic root, string taskName)
    {
        try
        {
            root.DeleteTask(taskName, 0);
        }
        catch
        {
            // Deleting an absent task is intentionally idempotent.
        }
    }

    private static void ReleaseComObject(object? value)
    {
        if (value is not null && System.Runtime.InteropServices.Marshal.IsComObject(value))
        {
            System.Runtime.InteropServices.Marshal.FinalReleaseComObject(value);
        }
    }
}