1  #pragma once
  2  
  3  #include <Windows.h>
  4  
  5  #include <functional>
  6  #include <string>
  7  #include <variant>
  8  #include <vector>
  9  #include <optional>
 10  #include <cassert>
 11  #include <sstream>
 12  
 13  #include "../logger/logger.h"
 14  #include "../utils/winapi_error.h"
 15  #include "../version/version.h"
 16  
 17  namespace registry
 18  {
 19      namespace detail
 20      {
 21          struct on_exit
 22          {
 23              std::function<void()> f;
 24  
 25              on_exit(std::function<void()> f) :
 26                  f{ std::move(f) } {}
 27              ~on_exit() { f(); }
 28          };
 29  
 30          template<class... Ts>
 31          struct overloaded : Ts...
 32          {
 33              using Ts::operator()...;
 34          };
 35  
 36          template<class... Ts>
 37          overloaded(Ts...) -> overloaded<Ts...>;
 38  
 39          inline const wchar_t* getScopeName(HKEY scope)
 40          {
 41              if (scope == HKEY_LOCAL_MACHINE)
 42              {
 43                  return L"HKLM";
 44              }
 45              else if (scope == HKEY_CURRENT_USER)
 46              {
 47                  return L"HKCU";
 48              }
 49              else if (scope == HKEY_CLASSES_ROOT)
 50              {
 51                  return L"HKCR";
 52              }
 53              else
 54              {
 55                  return L"HK??";
 56              }
 57          }
 58      }
 59  
 60      namespace install_scope
 61      {
 62          const wchar_t INSTALL_SCOPE_REG_KEY[] = L"Software\\Classes\\powertoys\\";
 63          const wchar_t UNINSTALL_REG_KEY[] = L"Software\\Microsoft\\Windows\\CurrentVersion\\Uninstall";
 64          
 65          // Bundle UpgradeCode from PowerToys.wxs (with braces as stored in registry)
 66          const wchar_t BUNDLE_UPGRADE_CODE[] = L"{6341382D-C0A9-4238-9188-BE9607E3FAB2}";
 67  
 68          enum class InstallScope
 69          {
 70              PerMachine = 0,
 71              PerUser,
 72          };
 73  
 74          // Helper function to find PowerToys bundle in Windows Uninstall registry by BundleUpgradeCode
 75          inline bool find_powertoys_bundle_in_uninstall_registry(HKEY rootKey)
 76          {
 77              HKEY uninstallKey{};
 78              if (RegOpenKeyExW(rootKey, UNINSTALL_REG_KEY, 0, KEY_READ, &uninstallKey) != ERROR_SUCCESS)
 79              {
 80                  return false;
 81              }
 82              detail::on_exit closeUninstallKey{ [uninstallKey] { RegCloseKey(uninstallKey); } };
 83  
 84              DWORD index = 0;
 85              wchar_t subKeyName[256];
 86  
 87              // Enumerate all subkeys under Uninstall
 88              while (RegEnumKeyW(uninstallKey, index++, subKeyName, 256) == ERROR_SUCCESS)
 89              {
 90                  HKEY productKey{};
 91                  if (RegOpenKeyExW(uninstallKey, subKeyName, 0, KEY_READ, &productKey) != ERROR_SUCCESS)
 92                  {
 93                      continue;
 94                  }
 95                  detail::on_exit closeProductKey{ [productKey] { RegCloseKey(productKey); } };
 96  
 97                  // Check BundleUpgradeCode value (specific to WiX Bundle installations)
 98                  wchar_t bundleUpgradeCode[256]{};
 99                  DWORD bundleUpgradeCodeSize = sizeof(bundleUpgradeCode);
100  
101                  if (RegQueryValueExW(productKey, L"BundleUpgradeCode", nullptr, nullptr,
102                                      reinterpret_cast<LPBYTE>(bundleUpgradeCode), &bundleUpgradeCodeSize) == ERROR_SUCCESS)
103                  {
104                      if (_wcsicmp(bundleUpgradeCode, BUNDLE_UPGRADE_CODE) == 0)
105                      {
106                          return true;
107                      }
108                  }
109              }
110  
111              return false;
112          }
113  
114          inline const InstallScope get_current_install_scope()
115          {
116              // 1. Check HKCU Uninstall registry first (user-level bundle)
117              // Note: MSI components are always in HKLM regardless of install scope,
118              // but the Bundle entry will be in HKCU for per-user installations
119              if (find_powertoys_bundle_in_uninstall_registry(HKEY_CURRENT_USER))
120              {
121                  Logger::info(L"Found user-level PowerToys bundle via BundleUpgradeCode in HKCU");
122                  return InstallScope::PerUser;
123              }
124  
125              // 2. Check HKLM Uninstall registry (machine-level bundle)
126              if (find_powertoys_bundle_in_uninstall_registry(HKEY_LOCAL_MACHINE))
127              {
128                  Logger::info(L"Found machine-level PowerToys bundle via BundleUpgradeCode in HKLM");
129                  return InstallScope::PerMachine;
130              }
131  
132              // 3. Fallback to legacy custom registry key detection
133              Logger::info(L"PowerToys bundle not found in Uninstall registry, falling back to legacy detection");
134  
135              // Open HKLM key
136              HKEY perMachineKey{};
137              if (RegOpenKeyExW(HKEY_LOCAL_MACHINE,
138                                INSTALL_SCOPE_REG_KEY,
139                                0,
140                                KEY_READ,
141                                &perMachineKey) != ERROR_SUCCESS)
142              {
143                  // Open HKCU key
144                  HKEY perUserKey{};
145                  if (RegOpenKeyExW(HKEY_CURRENT_USER,
146                                    INSTALL_SCOPE_REG_KEY,
147                                    0,
148                                    KEY_READ,
149                                    &perUserKey) != ERROR_SUCCESS)
150                  {
151                      // both keys are missing
152                      Logger::warn(L"No PowerToys installation detected, defaulting to PerMachine");
153                      return InstallScope::PerMachine;
154                  }
155                  else
156                  {
157                      DWORD dataSize{};
158                      if (RegGetValueW(
159                          perUserKey,
160                          nullptr,
161                          L"InstallScope",
162                          RRF_RT_REG_SZ,
163                          nullptr,
164                          nullptr,
165                          &dataSize) != ERROR_SUCCESS)
166                      {
167                          // HKCU key is missing
168                          RegCloseKey(perUserKey);
169                          return InstallScope::PerMachine;
170                      }
171  
172                      std::wstring data;
173                      data.resize(dataSize / sizeof(wchar_t));
174  
175                      if (RegGetValueW(
176                              perUserKey,
177                              nullptr,
178                              L"InstallScope",
179                              RRF_RT_REG_SZ,
180                              nullptr,
181                              &data[0],
182                              &dataSize) != ERROR_SUCCESS)
183                      {
184                          // HKCU key is missing
185                          RegCloseKey(perUserKey);
186                          return InstallScope::PerMachine;
187                      }
188                      RegCloseKey(perUserKey);
189  
190                      if (data.contains(L"perUser"))
191                      {
192                          return InstallScope::PerUser;
193                      }
194                  }
195              }
196  
197              return InstallScope::PerMachine;
198          }
199      }
200  
201      template<class>
202      inline constexpr bool always_false_v = false;
203  
204      struct ValueChange
205      {
206          using value_t = std::variant<DWORD, std::wstring>;
207          static constexpr size_t VALUE_BUFFER_SIZE = 512;
208  
209          HKEY scope{};
210          std::wstring path;
211          std::optional<std::wstring> name; // none == default
212          value_t value;
213          bool required = true;
214  
215          ValueChange(const HKEY scope, std::wstring path, std::optional<std::wstring> name, value_t value, bool required = true) :
216              scope{ scope }, path{ std::move(path) }, name{ std::move(name) }, value{ std::move(value) }, required{ required }
217          {
218          }
219  
220          std::wstring toString() const
221          {
222              using namespace detail;
223  
224              std::wstring value_str;
225              std::visit(overloaded{ [&](DWORD value) {
226                                        std::wostringstream oss;
227                                        oss << value;
228                                        value_str = oss.str();
229                                    },
230                                     [&](const std::wstring& value) { value_str = value; } },
231                         value);
232  
233              return fmt::format(L"{}\\{}\\{}:{}", detail::getScopeName(scope), path, name ? *name : L"Default", value_str);
234          }
235  
236          bool isApplied() const
237          {
238              HKEY key{};
239              if (auto res = RegOpenKeyExW(scope, path.c_str(), 0, KEY_READ, &key); res != ERROR_SUCCESS)
240              {
241                  Logger::info(L"isApplied of {}: RegOpenKeyExW failed: {}", toString(), get_last_error_or_default(res));
242                  return false;
243              }
244              detail::on_exit closeKey{ [key] { RegCloseKey(key); } };
245  
246              const DWORD expectedType = valueTypeToWinapiType(value);
247  
248              DWORD retrievedType{};
249              wchar_t buffer[VALUE_BUFFER_SIZE];
250              DWORD valueSize = sizeof(buffer);
251              if (auto res = RegQueryValueExW(key,
252                                              name.has_value() ? name->c_str() : nullptr,
253                                              0,
254                                              &retrievedType,
255                                              reinterpret_cast<LPBYTE>(&buffer),
256                                              &valueSize);
257                  res != ERROR_SUCCESS)
258              {
259                  Logger::info(L"isApplied of {}: RegQueryValueExW failed: {}", toString(), get_last_error_or_default(res));
260                  return false;
261              }
262  
263              if (expectedType != retrievedType)
264              {
265                  return false;
266              }
267  
268              if (const auto retrievedValue = bufferToValue(buffer, valueSize, retrievedType))
269              {
270                  return value == retrievedValue;
271              }
272              else
273              {
274                  return false;
275              }
276          }
277  
278          bool apply() const
279          {
280              HKEY key{};
281  
282              if (auto res = RegCreateKeyExW(scope, path.c_str(), 0, nullptr, REG_OPTION_NON_VOLATILE, KEY_WRITE, nullptr, &key, nullptr); res !=
283                                                                                                                                           ERROR_SUCCESS)
284              {
285                  Logger::error(L"apply of {}: RegCreateKeyExW failed: {}", toString(), get_last_error_or_default(res));
286                  return false;
287              }
288              detail::on_exit closeKey{ [key] { RegCloseKey(key); } };
289  
290              wchar_t buffer[VALUE_BUFFER_SIZE];
291              DWORD valueSize;
292              DWORD valueType;
293  
294              valueToBuffer(value, buffer, valueSize, valueType);
295              if (auto res = RegSetValueExW(key,
296                                            name.has_value() ? name->c_str() : nullptr,
297                                            0,
298                                            valueType,
299                                            reinterpret_cast<BYTE*>(buffer),
300                                            valueSize);
301                  res != ERROR_SUCCESS)
302              {
303                  Logger::error(L"apply of {}: RegSetValueExW failed: {}", toString(), get_last_error_or_default(res));
304                  return false;
305              }
306  
307              return true;
308          }
309  
310          bool unApply() const
311          {
312              HKEY key{};
313              if (auto res = RegOpenKeyExW(scope, path.c_str(), 0, KEY_ALL_ACCESS, &key); res != ERROR_SUCCESS)
314              {
315                  Logger::error(L"unApply of {}: RegOpenKeyExW failed: {}", toString(), get_last_error_or_default(res));
316                  return false;
317              }
318              detail::on_exit closeKey{ [key] { RegCloseKey(key); } };
319  
320              // delete the value itself
321              if (auto res = RegDeleteKeyValueW(scope, path.c_str(), name.has_value() ? name->c_str() : nullptr); res != ERROR_SUCCESS)
322              {
323                  Logger::error(L"unApply of {}: RegDeleteKeyValueW failed: {}", toString(), get_last_error_or_default(res));
324                  return false;
325              }
326  
327              // Check if the path doesn't contain anything and delete it if so
328              DWORD nValues = 0;
329              DWORD maxValueLen = 0;
330              const auto ok =
331                  RegQueryInfoKeyW(
332                      key, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, &nValues, nullptr, &maxValueLen, nullptr, nullptr) ==
333                  ERROR_SUCCESS;
334  
335              if (ok && (!nValues || !maxValueLen))
336              {
337                  RegDeleteTreeW(scope, path.c_str());
338              }
339              return true;
340          }
341  
342          bool requiresElevation() const { return scope == HKEY_LOCAL_MACHINE; }
343  
344      private:
345          static DWORD valueTypeToWinapiType(const value_t& v)
346          {
347              return std::visit(
348                  [](auto&& arg) {
349                      using T = std::decay_t<decltype(arg)>;
350                      if constexpr (std::is_same_v<T, DWORD>)
351                          return REG_DWORD;
352                      else if constexpr (std::is_same_v<T, std::wstring>)
353                          return REG_SZ;
354                      else
355                          static_assert(always_false_v<T>, "support for this registry type is not implemented");
356                  },
357                  v);
358          }
359  
360          static void valueToBuffer(const value_t& value, wchar_t buffer[VALUE_BUFFER_SIZE], DWORD& valueSize, DWORD& type)
361          {
362              using detail::overloaded;
363  
364              std::visit(overloaded{ [&](DWORD value) {
365                                        *reinterpret_cast<DWORD*>(buffer) = value;
366                                        type = REG_DWORD;
367                                        valueSize = sizeof(value);
368                                    },
369                                     [&](const std::wstring& value) {
370                                         assert(value.size() < VALUE_BUFFER_SIZE);
371                                         value.copy(buffer, value.size());
372                                         type = REG_SZ;
373                                         valueSize = static_cast<DWORD>(sizeof(wchar_t) * value.size());
374                                     } },
375                         value);
376          }
377  
378          static std::optional<value_t> bufferToValue(const wchar_t buffer[VALUE_BUFFER_SIZE],
379                                                      const DWORD valueSize,
380                                                      const DWORD type)
381          {
382              switch (type)
383              {
384              case REG_DWORD:
385                  return *reinterpret_cast<const DWORD*>(buffer);
386              case REG_SZ:
387              {
388                  if (!valueSize)
389                  {
390                      return std::wstring{};
391                  }
392                  std::wstring result{ buffer, valueSize / sizeof(wchar_t) };
393                  while (result[result.size() - 1] == L'\0')
394                  {
395                      result.resize(result.size() - 1);
396                  }
397                  return result;
398              }
399              default:
400                  return std::nullopt;
401              }
402          }
403      };
404  
405      struct ChangeSet
406      {
407          std::vector<ValueChange> changes;
408  
409          bool isApplied() const
410          {
411              for (const auto& c : changes)
412              {
413                  if (c.required && !c.isApplied())
414                  {
415                      return false;
416                  }
417              }
418              return true;
419          }
420  
421          bool apply() const
422          {
423              bool ok = true;
424              for (const auto& c : changes)
425              {
426                  ok = (c.apply()||!c.required) && ok;
427              }
428              return ok;
429          }
430  
431          bool unApply() const
432          {
433              bool ok = true;
434              for (const auto& c : changes)
435              {
436                  ok = (c.unApply()||!c.required) && ok;
437              }
438              return ok;
439          }
440      };
441  
442      const inline std::wstring DOTNET_COMPONENT_CATEGORY_CLSID = L"{62C8FE65-4EBB-45E7-B440-6E39B2CDBF29}";
443      const inline std::wstring ITHUMBNAIL_PROVIDER_CLSID = L"{E357FCCD-A995-4576-B01F-234630154E96}";
444      const inline std::wstring IPREVIEW_HANDLER_CLSID = L"{8895b1c6-b41f-4c1c-a562-0d564250836f}";
445  
446      namespace shellex
447      {
448          enum PreviewHandlerType
449          {
450              preview,
451              thumbnail
452          };
453  
454          inline registry::ChangeSet generatePreviewHandler(const PreviewHandlerType handlerType,
455                                                            const bool perUser,
456                                                            std::wstring handlerClsid,
457                                                            std::wstring powertoysVersion,
458                                                            std::wstring fullPathToHandler,
459                                                            std::wstring className,
460                                                            std::wstring displayName,
461                                                            std::vector<std::wstring> fileTypes,
462                                                            std::wstring perceivedType = L"",
463                                                            std::wstring fileKindType = L"")
464          {
465              const HKEY scope = perUser ? HKEY_CURRENT_USER : HKEY_LOCAL_MACHINE;
466  
467              std::wstring clsidPath = L"Software\\Classes\\CLSID";
468              clsidPath += L'\\';
469              clsidPath += handlerClsid;
470  
471              std::wstring inprocServerPath = clsidPath;
472              inprocServerPath += L'\\';
473              inprocServerPath += L"InprocServer32";
474  
475              std::wstring assemblyKeyValue;
476              if (const auto lastDotPos = className.rfind(L'.'); lastDotPos != std::wstring::npos)
477              {
478                  assemblyKeyValue = L"PowerToys." + className.substr(lastDotPos + 1);
479              }
480              else
481              {
482                  assemblyKeyValue = L"PowerToys." + className;
483              }
484  
485              assemblyKeyValue += L", Version=";
486              assemblyKeyValue += powertoysVersion;
487              assemblyKeyValue += L", Culture=neutral";
488  
489              std::wstring versionPath = inprocServerPath;
490              versionPath += L'\\';
491              versionPath += powertoysVersion;
492  
493              using vec_t = std::vector<registry::ValueChange>;
494              // TODO: verify that we actually need all of those
495              vec_t changes = { { scope, clsidPath, L"DisplayName", displayName },
496                                { scope, clsidPath, std::nullopt, className },
497                                { scope, inprocServerPath, std::nullopt, fullPathToHandler },
498                                { scope, inprocServerPath, L"Assembly", assemblyKeyValue },
499                                { scope, inprocServerPath, L"Class", className },
500                                { scope, inprocServerPath, L"ThreadingModel", L"Apartment" } };
501  
502              for (const auto& fileType : fileTypes)
503              {
504                  std::wstring fileTypePath = L"Software\\Classes\\" + fileType;
505                  std::wstring fileAssociationPath = fileTypePath + L"\\shellex\\";
506                  fileAssociationPath += handlerType == PreviewHandlerType::preview ? IPREVIEW_HANDLER_CLSID : ITHUMBNAIL_PROVIDER_CLSID;
507                  changes.push_back({ scope, fileAssociationPath, std::nullopt, handlerClsid });
508                  if (!fileKindType.empty())
509                  {
510                      // Registering a file type as a kind needs to be done at the HKEY_LOCAL_MACHINE level.
511                      // Make it optional as well so that we don't fail registering the handler if we can't write to HKEY_LOCAL_MACHINE.
512                      std::wstring kindMapPath = L"Software\\Microsoft\\Windows\\CurrentVersion\\Explorer\\KindMap";
513                      changes.push_back({ HKEY_LOCAL_MACHINE, kindMapPath, fileType, fileKindType, false});
514                  }
515                  if (!perceivedType.empty())
516                  {
517                      changes.push_back({ scope, fileTypePath, L"PerceivedType", perceivedType });
518                  }
519                  if (handlerType == PreviewHandlerType::preview && fileType == L".reg")
520                  {
521                      // this regfile registry key has precedence over Software\Classes\.reg for .reg files
522                      std::wstring regfilePath = L"Software\\Classes\\regfile\\shellex\\" + IPREVIEW_HANDLER_CLSID + L"\\";
523                      changes.push_back({ scope, regfilePath, std::nullopt, handlerClsid });
524                  }
525              }
526  
527              if (handlerType == PreviewHandlerType::preview)
528              {
529                  const std::wstring previewHostClsid = L"{6d2b5079-2f0b-48dd-ab7f-97cec514d30b}";
530                  const std::wstring previewHandlerListPath = LR"(Software\Microsoft\Windows\CurrentVersion\PreviewHandlers)";
531  
532                  changes.push_back({ scope, clsidPath, L"AppID", previewHostClsid });
533                  changes.push_back({ scope, previewHandlerListPath, handlerClsid, displayName });
534              }
535  
536              return registry::ChangeSet{ .changes = std::move(changes) };
537          }
538      }
539  }