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.Collections.Generic;
  7  using System.Collections.ObjectModel;
  8  using System.Diagnostics;
  9  using System.Linq;
 10  using System.Threading;
 11  using System.Threading.Tasks;
 12  
 13  using CommunityToolkit.Mvvm.ComponentModel;
 14  using CommunityToolkit.Mvvm.Input;
 15  using ManagedCommon;
 16  using PowerToys.FileLocksmithLib.Interop;
 17  
 18  namespace PowerToys.FileLocksmithUI.ViewModels
 19  {
 20  #pragma warning disable CA1708 // Identifiers should differ by more than case
 21      public partial class MainViewModel : ObservableObject, IDisposable
 22  #pragma warning restore CA1708 // Identifiers should differ by more than case
 23      {
 24          public IAsyncRelayCommand LoadProcessesCommand { get; }
 25  
 26          private bool _isLoading;
 27          private bool _isElevated;
 28          private string[] paths;
 29          private bool _disposed;
 30          private CancellationTokenSource _cancelProcessWatching;
 31  
 32          public ObservableCollection<ProcessResult> Processes { get; } = new();
 33  
 34          public bool IsLoading
 35          {
 36              get
 37              {
 38                  return _isLoading;
 39              }
 40  
 41              set
 42              {
 43                  _isLoading = value;
 44                  OnPropertyChanged(nameof(IsLoading));
 45              }
 46          }
 47  
 48          public bool IsElevated
 49          {
 50              get
 51              {
 52                  return _isElevated;
 53              }
 54  
 55              set
 56              {
 57                  _isElevated = value;
 58                  OnPropertyChanged(nameof(IsElevated));
 59              }
 60          }
 61  
 62          public string[] Paths
 63          {
 64              get => paths;
 65              set
 66              {
 67                  paths = value;
 68                  OnPropertyChanged(nameof(Paths));
 69              }
 70          }
 71  
 72          public string PathsToString
 73          {
 74              get
 75              {
 76                  return string.Join("\n", paths);
 77              }
 78          }
 79  
 80          public MainViewModel()
 81          {
 82              paths = NativeMethods.ReadPathsFromFile();
 83              Logger.LogInfo($"Starting FileLocksmith with {paths.Length} files selected.");
 84              LoadProcessesCommand = new AsyncRelayCommand(LoadProcessesAsync);
 85          }
 86  
 87          private async Task LoadProcessesAsync()
 88          {
 89              IsLoading = true;
 90              Processes.Clear();
 91  
 92              if (_cancelProcessWatching is not null)
 93              {
 94                  _cancelProcessWatching.Cancel();
 95              }
 96  
 97              _cancelProcessWatching = new CancellationTokenSource();
 98  
 99              var processes_found = await FindProcesses(paths);
100              if (processes_found is not null)
101              {
102                  foreach (ProcessResult p in processes_found)
103                  {
104                      Processes.Add(p);
105                      WatchProcess(p, _cancelProcessWatching.Token);
106                  }
107              }
108  
109              IsLoading = false;
110          }
111  
112          private async Task<List<ProcessResult>> FindProcesses(string[] paths)
113          {
114              var results = new List<ProcessResult>();
115              await Task.Run(() =>
116              {
117                  results = NativeMethods.FindProcessesRecursive(paths)?.ToList();
118              });
119              return results;
120          }
121  
122          private async void WatchProcess(ProcessResult process, CancellationToken token)
123          {
124              try
125              {
126                  Process handle = Process.GetProcessById((int)process.pid);
127                  try
128                  {
129                      await handle.WaitForExitAsync(token);
130                  }
131                  catch (TaskCanceledException)
132                  {
133                      // Nothing to do, normal operation
134                  }
135  
136                  if (handle.HasExited)
137                  {
138                      Processes.Remove(process);
139                  }
140              }
141              catch (Exception ex)
142              {
143                  Logger.LogError($"Couldn't add a waiter to wait for a process to exit. PID = {process.pid} and Name = {process.name}.", ex);
144                  Processes.Remove(process); // If we couldn't get a handle to the process or it has exited in the meanwhile, don't show it.
145              }
146          }
147  
148          [RelayCommand]
149          public void EndTask(ProcessResult selectedProcess)
150          {
151              try
152              {
153                  Process handle = Process.GetProcessById((int)selectedProcess.pid);
154                  try
155                  {
156                      handle.Kill();
157                  }
158                  catch (Exception ex)
159                  {
160                      Logger.LogError($"Couldn't kill process {selectedProcess.name} with PID {selectedProcess.pid}.", ex);
161                  }
162              }
163              catch (Exception ex)
164              {
165                  Logger.LogError($"Couldn't get a handle to kill process {selectedProcess.name} with PID {selectedProcess.pid}. Likely has been killed already.", ex);
166                  Processes.Remove(selectedProcess); // If we couldn't get a handle to the process, remove it from the list, since it's likely been killed already.
167              }
168          }
169  
170          [RelayCommand]
171          public void RestartElevated()
172          {
173              if (NativeMethods.StartAsElevated(paths))
174              {
175                  // TODO gentler exit
176                  Environment.Exit(0);
177              }
178              else
179              {
180                  // TODO report error?
181                  Logger.LogError($"Couldn't restart as elevated.");
182              }
183          }
184  
185          public void Dispose()
186          {
187              Dispose(disposing: true);
188              GC.SuppressFinalize(this);
189          }
190  
191          protected virtual void Dispose(bool disposing)
192          {
193              if (!_disposed)
194              {
195                  if (disposing)
196                  {
197                      _disposed = true;
198                  }
199              }
200          }
201      }
202  }