LenovoLegionToolkit

1
using System;
2
using System.Collections.Generic;
3
using System.ComponentModel;
4
using System.Diagnostics;
5
using System.Linq;
6
using System.ServiceProcess;
7
using System.Threading.Tasks;
8
using LenovoLegionToolkit.Lib.Extensions;
9
using LenovoLegionToolkit.Lib.Utils;
10
using TaskService = Microsoft.Win32.TaskScheduler.TaskService;
11

12
namespace LenovoLegionToolkit.Lib.SoftwareDisabler;
13

14
public class SoftwareDisablerException : Exception
15
{
16
    public SoftwareDisablerException(string message, Exception innerException) : base(message, innerException) { }
17
}
18

19
public abstract class AbstractSoftwareDisabler
20
{
21
    protected abstract IEnumerable<string> ScheduledTasksPaths { get; }
22
    protected abstract IEnumerable<string> ServiceNames { get; }
23
    protected abstract IEnumerable<string> ProcessNames { get; }
24

25
    public Task<SoftwareStatus> GetStatusAsync() => Task.Run(() =>
26
    {
27
        bool isEnabled;
28
        bool isInstalled;
29

30
        try
31
        {
32
            var services = RunningServices().ToArray();
33
            var processes = RunningProcesses().ToArray();
34

35
            if (Log.Instance.IsTraceEnabled)
36
            {
37
                Log.Instance.Trace($"Running services count: {services.Length}. [type={GetType().Name}, services={string.Join(",", services)}]");
38
                Log.Instance.Trace($"Running processes count: {processes.Length}. [type={GetType().Name}, processes={string.Join(",", processes)}]");
39
            }
40

41
            isEnabled = services.Any() || processes.Any();
42
            isInstalled = IsInstalled();
43
        }
44
        catch (Exception ex)
45
        {
46
            if (Log.Instance.IsTraceEnabled)
47
                Log.Instance.Trace($"Exception while getting status. [type={GetType().Name}]", ex);
48

49
            isEnabled = false;
50
            isInstalled = false;
51
        }
52

53
        if (Log.Instance.IsTraceEnabled)
54
            Log.Instance.Trace($"Status: {isEnabled},{isInstalled} [type={GetType().Name}]");
55

56
        if (isEnabled)
57
            return SoftwareStatus.Enabled;
58

59
        if (!isInstalled)
60
            return SoftwareStatus.NotFound;
61

62
        return SoftwareStatus.Disabled;
63
    });
64

65
    public virtual Task EnableAsync() => Task.Run(() =>
66
    {
67
        if (Log.Instance.IsTraceEnabled)
68
            Log.Instance.Trace($"Enabling... [type={GetType().Name}]");
69

70
        SetScheduledTasksEnabled(true);
71
        SetServicesEnabled(true);
72

73
        if (Log.Instance.IsTraceEnabled)
74
            Log.Instance.Trace($"Enabled [type={GetType().Name}]");
75
    });
76

77
    public virtual Task DisableAsync() => Task.Run(async () =>
78
    {
79
        if (Log.Instance.IsTraceEnabled)
80
            Log.Instance.Trace($"Disabling... [type={GetType().Name}]");
81

82
        SetScheduledTasksEnabled(false);
83
        SetServicesEnabled(false);
84
        await KillProcessesAsync().ConfigureAwait(false);
85

86
        if (Log.Instance.IsTraceEnabled)
87
            Log.Instance.Trace($"Disabled [type={GetType().Name}]");
88
    });
89

90
    private bool IsInstalled() => ServiceController.GetServices().Any(s => ServiceNames.Contains(s.ServiceName));
91

92
    private IEnumerable<string> RunningServices()
93
    {
94
        var services = ServiceController.GetServices();
95
        return ServiceNames.Where(s => IsServiceEnabled(s, services));
96
    }
97

98
    protected virtual IEnumerable<string> RunningProcesses()
99
    {
100
        foreach (var process in Process.GetProcesses())
101
        {
102
            foreach (var processName in ProcessNames)
103
            {
104
                var name = string.Empty;
105

106
                try
107
                {
108
                    name = process.ProcessName;
109
                    if (!name.StartsWith(processName, StringComparison.InvariantCultureIgnoreCase))
110
                        continue;
111
                }
112
                catch {  /* Ignored. */ }
113

114
                if (!string.IsNullOrEmpty(name))
115
                    yield return name;
116
            }
117
        }
118
    }
119

120
    private static bool IsServiceEnabled(string serviceName, IEnumerable<ServiceController> services)
121
    {
122
        try
123
        {
124
            var service = services.FirstOrDefault(s => s.ServiceName == serviceName);
125
            if (service is null)
126
                return false;
127

128
            return service.Status is not ServiceControllerStatus.Stopped;
129
        }
130
        catch (InvalidOperationException)
131
        {
132
            return false;
133
        }
134
    }
135

136
    private void SetScheduledTasksEnabled(bool enabled)
137
    {
138
        var taskService = TaskService.Instance;
139
        foreach (var path in ScheduledTasksPaths)
140
            SetTasksInFolderEnabled(taskService, path, enabled);
141
    }
142

143
    private void SetTasksInFolderEnabled(TaskService taskService, string path, bool enabled)
144
    {
145
        if (Log.Instance.IsTraceEnabled)
146
            Log.Instance.Trace($"Setting tasks in folder {path} to {enabled}. [type={GetType().Name}]");
147

148
        var folder = taskService.GetFolder(path);
149
        if (folder is null)
150
        {
151
            if (Log.Instance.IsTraceEnabled)
152
                Log.Instance.Trace($"Folder not found [path={path}, type={GetType().Name}]]");
153

154
            return;
155
        }
156

157
        foreach (var task in folder.Tasks.ToArray())
158
        {
159
            task.Definition.Settings.Enabled = enabled;
160
            try
161
            {
162
                task.RegisterChanges();
163
            }
164
            catch (Exception ex)
165
            {
166
                if (Log.Instance.IsTraceEnabled)
167
                    Log.Instance.Trace($"Failed to register changes on task {task.Name} in {task.Path}.", ex);
168

169
                throw new SoftwareDisablerException($"Failed to register changes on task {task.Name} in {task.Path}. [type={GetType().Name}]", ex);
170
            }
171
        }
172
    }
173

174
    private void SetServicesEnabled(bool enabled)
175
    {
176
        foreach (var serviceName in ServiceNames)
177
            SetServiceEnabled(serviceName, enabled);
178
    }
179

180
    private void SetServiceEnabled(string serviceName, bool enabled)
181
    {
182
        try
183
        {
184
            if (Log.Instance.IsTraceEnabled)
185
                Log.Instance.Trace($"Setting service {serviceName} to {enabled}. [type={GetType().Name}]");
186

187
            if (!ServiceController.GetServices().Any(s => s.ServiceName == serviceName))
188
            {
189
                if (Log.Instance.IsTraceEnabled)
190
                    Log.Instance.Trace($"Service {serviceName} not found. [type={GetType().Name}]");
191

192
                return;
193
            }
194

195
            var service = new ServiceController(serviceName);
196

197
            try
198
            {
199
                if (Log.Instance.IsTraceEnabled)
200
                    Log.Instance.Trace($"Changing service {serviceName} start mode to {enabled}.  [type={GetType().Name}]");
201

202
                service.ChangeStartMode(enabled);
203

204
                if (enabled)
205
                {
206
                    if (service.Status != ServiceControllerStatus.Running)
207
                    {
208
                        if (Log.Instance.IsTraceEnabled)
209
                            Log.Instance.Trace($"Starting service {serviceName}... [type={GetType().Name}]");
210
                        service.Start();
211
                        service.WaitForStatus(ServiceControllerStatus.Running);
212
                    }
213
                    else
214
                    {
215
                        if (Log.Instance.IsTraceEnabled)
216
                            Log.Instance.Trace($"Will not start service {serviceName}. [status={service.Status}, type={GetType().Name}]]");
217
                    }
218
                }
219
                else
220
                {
221
                    if (service.CanStop)
222
                    {
223
                        if (Log.Instance.IsTraceEnabled)
224
                            Log.Instance.Trace($"Stopping service {serviceName}... [type={GetType().Name}]");
225
                        service.Stop();
226
                        service.WaitForStatus(ServiceControllerStatus.Stopped);
227
                    }
228
                    else
229
                    {
230
                        if (Log.Instance.IsTraceEnabled)
231
                            Log.Instance.Trace($"Will not stop service {serviceName}. [status={service.Status}, canStop={service.CanStop}, type={GetType().Name}]]");
232
                    }
233
                }
234
            }
235
            finally
236
            {
237
                service.Close();
238
            }
239
        }
240
        catch (InvalidOperationException ex) when (ex.InnerException is Win32Exception { NativeErrorCode: 1060 })
241
        {
242
            if (Log.Instance.IsTraceEnabled)
243
                Log.Instance.Trace($"Service {serviceName} could not be set to {enabled}");
244

245
            throw new SoftwareDisablerException(serviceName, ex);
246
        }
247
    }
248

249
    protected virtual async Task KillProcessesAsync()
250
    {
251
        foreach (var process in Process.GetProcesses())
252
            foreach (var processName in ProcessNames)
253
            {
254
                try
255
                {
256
                    if (process.ProcessName.StartsWith(processName, StringComparison.InvariantCultureIgnoreCase))
257
                    {
258
                        process.Kill();
259
                        await process.WaitForExitAsync().ConfigureAwait(false);
260
                    }
261
                }
262
                catch (Exception ex)
263
                {
264
                    if (Log.Instance.IsTraceEnabled)
265
                        Log.Instance.Trace($"Couldn't kill process.", ex);
266
                }
267
            }
268
    }
269
}
270