1  // Copyright (c) Microsoft Corporation
  2  // The Microsoft Corporation licenses this file to you under the MIT license.
  3  // See the LICENSE file in the project root for more information.
  4  
  5  using System;
  6  using System.CommandLine;
  7  using System.CommandLine.Parsing;
  8  using System.Diagnostics;
  9  using System.Globalization;
 10  using System.IO;
 11  using System.Linq;
 12  using System.Reactive.Concurrency;
 13  using System.Reactive.Linq;
 14  using System.Reflection;
 15  using System.Text.Json;
 16  using System.Threading;
 17  using System.Threading.Tasks;
 18  using Awake.Core;
 19  using Awake.Core.Models;
 20  using Awake.Core.Native;
 21  using Awake.Properties;
 22  using ManagedCommon;
 23  using Microsoft.PowerToys.Settings.UI.Library;
 24  using Microsoft.PowerToys.Telemetry;
 25  
 26  namespace Awake
 27  {
 28      internal sealed class Program
 29      {
 30          private static readonly string[] _aliasesConfigOption = ["--use-pt-config", "-c"];
 31          private static readonly string[] _aliasesDisplayOption = ["--display-on", "-d"];
 32          private static readonly string[] _aliasesTimeOption = ["--time-limit", "-t"];
 33          private static readonly string[] _aliasesPidOption = ["--pid", "-p"];
 34          private static readonly string[] _aliasesExpireAtOption = ["--expire-at", "-e"];
 35          private static readonly string[] _aliasesParentPidOption = ["--use-parent-pid", "-u"];
 36  
 37          private static readonly JsonSerializerOptions _serializerOptions = new() { IncludeFields = true };
 38          private static readonly ETWTrace _etwTrace = new();
 39  
 40          private static FileSystemWatcher? _watcher;
 41          private static SettingsUtils? _settingsUtils;
 42          private static EventWaitHandle? _exitEventHandle;
 43          private static RegisteredWaitHandle? _registeredWaitHandle;
 44  
 45          private static bool _startedFromPowerToys;
 46  
 47          public static Mutex? LockMutex { get; set; }
 48  
 49          private static ConsoleEventHandler? _handler;
 50          private static SystemPowerCapabilities _powerCapabilities;
 51  
 52          private static async Task<int> Main(string[] args)
 53          {
 54              Logger.InitializeLogger(Path.Combine("\\", Core.Constants.AppName, "Logs"));
 55  
 56              var rootCommand = BuildRootCommand();
 57  
 58              Bridge.AttachConsole(Core.Native.Constants.ATTACH_PARENT_PROCESS);
 59  
 60              var parseResult = rootCommand.Parse(args);
 61  
 62              if (parseResult.Tokens.Any(t => t.Value.ToLowerInvariant() is "--help" or "-h" or "-?"))
 63              {
 64                  // Print help and exit.
 65                  return rootCommand.Invoke(args);
 66              }
 67  
 68              if (parseResult.Errors.Count > 0)
 69              {
 70                  // Shows errors and returns non-zero.
 71                  return rootCommand.Invoke(args);
 72              }
 73  
 74              _settingsUtils = SettingsUtils.Default;
 75  
 76              LockMutex = new Mutex(true, Core.Constants.AppName, out bool instantiated);
 77  
 78              try
 79              {
 80                  string appLanguage = LanguageHelper.LoadLanguage();
 81                  if (!string.IsNullOrEmpty(appLanguage))
 82                  {
 83                      Thread.CurrentThread.CurrentUICulture = new CultureInfo(appLanguage);
 84                  }
 85              }
 86              catch (CultureNotFoundException ex)
 87              {
 88                  Logger.LogError("CultureNotFoundException: " + ex.Message);
 89              }
 90  
 91              await TrayHelper.InitializeTray(TrayHelper.DefaultAwakeIcon, Core.Constants.FullAppName);
 92              AppDomain.CurrentDomain.ProcessExit += (_, _) => TrayHelper.RunOnMainThread(() => LockMutex?.ReleaseMutex());
 93              AppDomain.CurrentDomain.UnhandledException += AwakeUnhandledExceptionCatcher;
 94  
 95              if (!instantiated)
 96              {
 97                  // Awake is already running - there is no need for us to process
 98                  // anything further
 99                  Exit(Core.Constants.AppName + " is already running! Exiting the application.", 1);
100                  return 1;
101              }
102              else
103              {
104                  if (PowerToys.GPOWrapper.GPOWrapper.GetConfiguredAwakeEnabledValue() == PowerToys.GPOWrapper.GpoRuleConfigured.Disabled)
105                  {
106                      Exit("PowerToys.Awake tried to start with a group policy setting that disables the tool. Please contact your system administrator.", 1);
107                      return 1;
108                  }
109                  else
110                  {
111                      Logger.LogInfo($"Launching {Core.Constants.AppName}...");
112                      Logger.LogInfo(FileVersionInfo.GetVersionInfo(Assembly.GetExecutingAssembly().Location).FileVersion);
113                      Logger.LogInfo($"Build: {Core.Constants.BuildId}");
114                      Logger.LogInfo($"OS: {Environment.OSVersion}");
115                      Logger.LogInfo($"OS Build: {Manager.GetOperatingSystemBuild()}");
116  
117                      TaskScheduler.UnobservedTaskException += (sender, args) =>
118                      {
119                          Trace.WriteLine($"Task scheduler error: {args.Exception.Message}"); // somebody forgot to check!
120                          args.SetObserved();
121                      };
122  
123                      // To make it easier to diagnose future issues, let's get the
124                      // system power capabilities and aggregate them in the log.
125                      Bridge.GetPwrCapabilities(out _powerCapabilities);
126                      Logger.LogInfo(JsonSerializer.Serialize(_powerCapabilities, _serializerOptions));
127  
128                      return await rootCommand.InvokeAsync(args);
129                  }
130              }
131          }
132  
133          private static RootCommand BuildRootCommand()
134          {
135              Logger.LogInfo("Parsing parameters...");
136  
137              Option<bool> configOption = new(_aliasesConfigOption, () => false, Resources.AWAKE_CMD_HELP_CONFIG_OPTION)
138              {
139                  Arity = ArgumentArity.ZeroOrOne,
140                  IsRequired = false,
141              };
142  
143              Option<bool> displayOption = new(_aliasesDisplayOption, () => false, Resources.AWAKE_CMD_HELP_DISPLAY_OPTION)
144              {
145                  Arity = ArgumentArity.ZeroOrOne,
146                  IsRequired = false,
147              };
148  
149              Option<uint> timeOption = new(_aliasesTimeOption, () => 0, Resources.AWAKE_CMD_HELP_TIME_OPTION)
150              {
151                  Arity = ArgumentArity.ExactlyOne,
152                  IsRequired = false,
153              };
154  
155              Option<int> pidOption = new(_aliasesPidOption, () => 0, Resources.AWAKE_CMD_HELP_PID_OPTION)
156              {
157                  Arity = ArgumentArity.ZeroOrOne,
158                  IsRequired = false,
159              };
160  
161              Option<string> expireAtOption = new(_aliasesExpireAtOption, () => string.Empty, Resources.AWAKE_CMD_HELP_EXPIRE_AT_OPTION)
162              {
163                  Arity = ArgumentArity.ZeroOrOne,
164                  IsRequired = false,
165              };
166  
167              Option<bool> parentPidOption = new(_aliasesParentPidOption, () => false, Resources.AWAKE_CMD_PARENT_PID_OPTION)
168              {
169                  Arity = ArgumentArity.ZeroOrOne,
170                  IsRequired = false,
171              };
172  
173              timeOption.AddValidator(result =>
174              {
175                  if (result.Tokens.Count != 0 && !uint.TryParse(result.Tokens[0].Value, out _))
176                  {
177                      string errorMessage = $"Interval in --time-limit could not be parsed correctly. Check that the value is valid and doesn't exceed 4,294,967,295. Value used: {result.Tokens[0].Value}.";
178                      Logger.LogError(errorMessage);
179                      result.ErrorMessage = errorMessage;
180                  }
181              });
182  
183              pidOption.AddValidator(result =>
184              {
185                  if (result.Tokens.Count != 0 && !int.TryParse(result.Tokens[0].Value, out _))
186                  {
187                      string errorMessage = $"PID value in --pid could not be parsed correctly. Check that the value is valid and falls within the boundaries of Windows PID process limits. Value used: {result.Tokens[0].Value}.";
188                      Logger.LogError(errorMessage);
189                      result.ErrorMessage = errorMessage;
190                  }
191              });
192  
193              expireAtOption.AddValidator(result =>
194              {
195                  if (result.Tokens.Count != 0 && !DateTimeOffset.TryParse(result.Tokens[0].Value, out _))
196                  {
197                      string errorMessage = $"Date and time value in --expire-at could not be parsed correctly. Check that the value is valid date and time. Refer to https://aka.ms/powertoys/awake for format examples. Value used: {result.Tokens[0].Value}.";
198                      Logger.LogError(errorMessage);
199                      result.ErrorMessage = errorMessage;
200                  }
201              });
202  
203              RootCommand? rootCommand =
204              [
205                  configOption,
206                  displayOption,
207                  timeOption,
208                  pidOption,
209                  expireAtOption,
210                  parentPidOption,
211              ];
212  
213              rootCommand.Description = Core.Constants.AppName;
214              rootCommand.SetHandler(HandleCommandLineArguments, configOption, displayOption, timeOption, pidOption, expireAtOption, parentPidOption);
215  
216              return rootCommand;
217          }
218  
219          private static void AwakeUnhandledExceptionCatcher(object sender, UnhandledExceptionEventArgs e)
220          {
221              if (e.ExceptionObject is Exception exception)
222              {
223                  Logger.LogError(exception.ToString());
224                  Logger.LogError(exception.StackTrace);
225              }
226          }
227  
228          private static bool ExitHandler(ControlType ctrlType)
229          {
230              Logger.LogInfo($"Exited through handler with control type: {ctrlType}");
231              Exit(Resources.AWAKE_EXIT_MESSAGE, Environment.ExitCode);
232              return false;
233          }
234  
235          private static void Exit(string message, int exitCode)
236          {
237              _etwTrace?.Dispose();
238              DisposeFileSystemWatcher();
239              _registeredWaitHandle?.Unregister(null);
240              _exitEventHandle?.Dispose();
241              Logger.LogInfo(message);
242              Manager.CompleteExit(exitCode);
243          }
244  
245          private static void DisposeFileSystemWatcher()
246          {
247              if (_watcher != null)
248              {
249                  _watcher.EnableRaisingEvents = false;
250                  _watcher.Dispose();
251                  _watcher = null;
252              }
253          }
254  
255          private static bool ProcessExists(int processId)
256          {
257              if (processId <= 0)
258              {
259                  return false;
260              }
261  
262              try
263              {
264                  // Throws if the Process ID is not found.
265                  using var p = Process.GetProcessById(processId);
266                  return !p.HasExited;
267              }
268              catch (ArgumentException)
269              {
270                  // Process with the specified ID is not running
271                  return false;
272              }
273              catch (InvalidOperationException ex)
274              {
275                  // Process has exited or cannot be accessed
276                  Logger.LogInfo($"Process {processId} cannot be accessed: {ex.Message}");
277                  return false;
278              }
279          }
280  
281          private static void HandleCommandLineArguments(bool usePtConfig, bool displayOn, uint timeLimit, int pid, string expireAt, bool useParentPid)
282          {
283              if (pid == 0 && !useParentPid)
284              {
285                  Logger.LogInfo("No PID specified. Allocating console...");
286                  Bridge.FreeConsole();
287                  AllocateLocalConsole();
288              }
289              else
290              {
291                  Logger.LogInfo("Starting with PID binding.");
292                  _startedFromPowerToys = true;
293              }
294  
295              Logger.LogInfo($"The value for --use-pt-config is: {usePtConfig}");
296              Logger.LogInfo($"The value for --display-on is: {displayOn}");
297              Logger.LogInfo($"The value for --time-limit is: {timeLimit}");
298              Logger.LogInfo($"The value for --pid is: {pid}");
299              Logger.LogInfo($"The value for --expire-at is: {expireAt}");
300              Logger.LogInfo($"The value for --use-parent-pid is: {useParentPid}");
301  
302              // Start the monitor thread that will be used to track the current state.
303              Manager.StartMonitor();
304  
305              _exitEventHandle = new EventWaitHandle(false, EventResetMode.ManualReset, PowerToys.Interop.Constants.AwakeExitEvent());
306              _registeredWaitHandle = ThreadPool.RegisterWaitForSingleObject(
307                  _exitEventHandle,
308                  (state, timedOut) => Exit(Resources.AWAKE_EXIT_SIGNAL_MESSAGE, 0),
309                  null,
310                  Timeout.Infinite,
311                  executeOnlyOnce: true);
312  
313              if (usePtConfig)
314              {
315                  // Configuration file is used, therefore we disregard any other command-line parameter
316                  // and instead watch for changes in the file. This is used as a priority against all other arguments,
317                  // so if --use-pt-config is applied the rest of the arguments are irrelevant.
318                  Manager.IsUsingPowerToysConfig = true;
319  
320                  try
321                  {
322                      string? settingsPath = _settingsUtils!.GetSettingsFilePath(Core.Constants.AppName);
323  
324                      Logger.LogInfo($"Reading configuration file: {settingsPath}");
325  
326                      if (!File.Exists(settingsPath))
327                      {
328                          Logger.LogError("The settings file does not exist. Scaffolding default configuration...");
329  
330                          AwakeSettings scaffoldSettings = new();
331                          _settingsUtils.SaveSettings(JsonSerializer.Serialize(scaffoldSettings), Core.Constants.AppName);
332                      }
333  
334                      ScaffoldConfiguration(settingsPath);
335  
336                      if (pid != 0)
337                      {
338                          if (!ProcessExists(pid))
339                          {
340                              Logger.LogError($"PID {pid} does not exist or is not accessible. Exiting.");
341                              Exit(Resources.AWAKE_EXIT_PROCESS_BINDING_FAILURE_MESSAGE, 1);
342                          }
343  
344                          Logger.LogInfo($"Bound to target process while also using PowerToys settings: {pid}");
345  
346                          RunnerHelper.WaitForPowerToysRunner(pid, () =>
347                          {
348                              Logger.LogInfo($"Triggered PID-based exit handler for PID {pid}.");
349                              Exit(Resources.AWAKE_EXIT_BINDING_HOOK_MESSAGE, 0);
350                          });
351                      }
352                  }
353                  catch (Exception ex)
354                  {
355                      Logger.LogError($"There was a problem with the configuration file. Make sure it exists. {ex.Message}");
356                  }
357              }
358              else if (pid != 0 || useParentPid)
359              {
360                  HandleProcessScopedKeepAwake(pid, useParentPid, displayOn);
361              }
362              else
363              {
364                  // Date-based binding takes precedence over timed configuration, so we want to
365                  // check for that first.
366                  if (!string.IsNullOrWhiteSpace(expireAt))
367                  {
368                      try
369                      {
370                          DateTimeOffset expirationDateTime = DateTimeOffset.Parse(expireAt, CultureInfo.CurrentCulture);
371                          Logger.LogInfo($"Operating in thread ID {Environment.CurrentManagedThreadId}.");
372                          Manager.SetExpirableKeepAwake(expirationDateTime, displayOn);
373                      }
374                      catch (Exception ex)
375                      {
376                          Logger.LogError($"Could not parse date string {expireAt} into a DateTimeOffset object.");
377                          Logger.LogError(ex.Message);
378                      }
379                  }
380                  else
381                  {
382                      AwakeMode mode = timeLimit <= 0 ? AwakeMode.INDEFINITE : AwakeMode.TIMED;
383  
384                      if (mode == AwakeMode.INDEFINITE)
385                      {
386                          Manager.SetIndefiniteKeepAwake(displayOn);
387                      }
388                      else
389                      {
390                          Manager.SetTimedKeepAwake(timeLimit, displayOn);
391                      }
392                  }
393              }
394          }
395  
396          /// <summary>
397          /// Start a process-scoped keep-awake session. The application will keep the system awake
398          /// indefinitely until the target process terminates.
399          /// </summary>
400          /// <param name="pid">The explicit process ID to monitor.</param>
401          /// <param name="useParentPid">A flag indicating whether the application should monitor its
402          /// parent process.</param>
403          /// <param name="displayOn">Whether to keep the display on during the session.</param>
404          private static void HandleProcessScopedKeepAwake(int pid, bool useParentPid, bool displayOn)
405          {
406              int targetPid = 0;
407  
408              // We prioritize a user-provided PID over the parent PID. If both are given on the
409              // command line, the --pid value will be used.
410              if (pid != 0)
411              {
412                  if (pid == Environment.ProcessId)
413                  {
414                      Logger.LogError("Awake cannot bind to itself, as this would lead to an indefinite keep-awake state.");
415                      Exit(Resources.AWAKE_EXIT_BIND_TO_SELF_FAILURE_MESSAGE, 1);
416                  }
417  
418                  if (!ProcessExists(pid))
419                  {
420                      Logger.LogError($"PID {pid} does not exist or is not accessible. Exiting.");
421                      Exit(Resources.AWAKE_EXIT_PROCESS_BINDING_FAILURE_MESSAGE, 1);
422                  }
423  
424                  targetPid = pid;
425              }
426              else if (useParentPid)
427              {
428                  targetPid = Manager.GetParentProcess()?.Id ?? 0;
429  
430                  if (targetPid == 0)
431                  {
432                      // The parent process could not be identified.
433                      Logger.LogError("Failed to identify a parent process for binding.");
434                      Exit(Resources.AWAKE_EXIT_PARENT_BINDING_FAILURE_MESSAGE, 1);
435                  }
436              }
437  
438              // We have a valid non-zero PID to monitor.
439              Logger.LogInfo($"Bound to target process: {targetPid}");
440  
441              // Sets the keep-awake plan and updates the tray icon.
442              Manager.SetIndefiniteKeepAwake(displayOn, targetPid);
443  
444              // Synchronize with the target process, and trigger Exit() when it finishes.
445              RunnerHelper.WaitForPowerToysRunner(targetPid, () =>
446              {
447                  Logger.LogInfo($"Triggered PID-based exit handler for PID {targetPid}.");
448                  Exit(Resources.AWAKE_EXIT_BINDING_HOOK_MESSAGE, 0);
449              });
450          }
451  
452          private static void AllocateLocalConsole()
453          {
454              Manager.AllocateConsole();
455  
456              _handler = new ConsoleEventHandler(ExitHandler);
457              Manager.SetConsoleControlHandler(_handler, true);
458  
459              Trace.Listeners.Add(new ConsoleTraceListener());
460          }
461  
462          private static void ScaffoldConfiguration(string settingsPath)
463          {
464              try
465              {
466                  SetupFileSystemWatcher(settingsPath);
467                  InitializeSettings();
468                  ProcessSettings();
469              }
470              catch (Exception ex)
471              {
472                  Logger.LogError($"An error occurred scaffolding the configuration. Error details: {ex.Message}");
473              }
474          }
475  
476          private static void SetupFileSystemWatcher(string settingsPath)
477          {
478              string directory = Path.GetDirectoryName(settingsPath)!;
479              string fileName = Path.GetFileName(settingsPath);
480  
481              _watcher = new FileSystemWatcher
482              {
483                  Path = directory,
484                  EnableRaisingEvents = true,
485                  NotifyFilter = NotifyFilters.LastWrite | NotifyFilters.CreationTime,
486                  Filter = fileName,
487              };
488  
489              Observable.Merge(
490                  Observable.FromEventPattern<FileSystemEventHandler, FileSystemEventArgs>(
491                      h => _watcher.Changed += h,
492                      h => _watcher.Changed -= h),
493                  Observable.FromEventPattern<FileSystemEventHandler, FileSystemEventArgs>(
494                      h => _watcher.Created += h,
495                      h => _watcher.Created -= h))
496              .Throttle(TimeSpan.FromMilliseconds(25))
497              .SubscribeOn(TaskPoolScheduler.Default)
498              .Select(e => e.EventArgs)
499              .Subscribe(HandleAwakeConfigChange);
500          }
501  
502          private static void InitializeSettings()
503          {
504              AwakeSettings settings = Manager.ModuleSettings?.GetSettings<AwakeSettings>(Core.Constants.AppName) ?? new AwakeSettings();
505              TrayHelper.SetTray(settings, _startedFromPowerToys);
506          }
507  
508          private static void HandleAwakeConfigChange(FileSystemEventArgs fileEvent)
509          {
510              try
511              {
512                  Logger.LogInfo("Detected a settings file change. Updating configuration...");
513                  ProcessSettings();
514              }
515              catch (Exception e)
516              {
517                  Logger.LogError($"Could not handle Awake configuration change. Error: {e.Message}");
518              }
519          }
520  
521          private static void ProcessSettings()
522          {
523              try
524              {
525                  AwakeSettings settings = _settingsUtils!.GetSettings<AwakeSettings>(Core.Constants.AppName)
526                      ?? throw new InvalidOperationException("Settings are null.");
527  
528                  Logger.LogInfo($"Identified custom time shortcuts for the tray: {settings.Properties.CustomTrayTimes.Count}");
529  
530                  switch (settings.Properties.Mode)
531                  {
532                      case AwakeMode.PASSIVE:
533                          Manager.SetPassiveKeepAwake();
534                          break;
535  
536                      case AwakeMode.INDEFINITE:
537                          Manager.SetIndefiniteKeepAwake(settings.Properties.KeepDisplayOn);
538                          break;
539  
540                      case AwakeMode.TIMED:
541                          uint computedTime = (settings.Properties.IntervalHours * 3600) + (settings.Properties.IntervalMinutes * 60);
542                          Manager.SetTimedKeepAwake(computedTime, settings.Properties.KeepDisplayOn);
543                          break;
544  
545                      case AwakeMode.EXPIRABLE:
546                          // When we are loading from the settings file, let's make sure that we never
547                          // get users in a state where the expirable keep-awake is in the past.
548                          if (settings.Properties.ExpirationDateTime <= DateTimeOffset.Now)
549                          {
550                              settings.Properties.ExpirationDateTime = DateTimeOffset.Now.AddMinutes(5);
551                              _settingsUtils.SaveSettings(JsonSerializer.Serialize(settings), Core.Constants.AppName);
552  
553                              // Return here - the FileSystemWatcher will re-trigger ProcessSettings
554                              // with the corrected expiration time, which will then call SetExpirableKeepAwake.
555                              // This matches the pattern used by mode setters (e.g., SetExpirableKeepAwake line 292).
556                              return;
557                          }
558  
559                          Manager.SetExpirableKeepAwake(settings.Properties.ExpirationDateTime, settings.Properties.KeepDisplayOn);
560                          break;
561  
562                      default:
563                          Logger.LogError("Unknown mode of operation. Check config file.");
564                          break;
565                  }
566  
567                  TrayHelper.SetTray(settings, _startedFromPowerToys);
568              }
569              catch (Exception ex)
570              {
571                  Logger.LogError($"There was a problem reading the configuration file. Error: {ex.GetType()} {ex.Message}");
572              }
573          }
574      }
575  }