-static BOOL set_privilege(LPCWSTR privilegeName, BOOL enabled)
-{
- HANDLE hToken = INVALID_HANDLE_VALUE;
- TOKEN_PRIVILEGES tp;
- DWORD error = ERROR_SUCCESS;
-
- if (OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES, &hToken))
- {
- tp.PrivilegeCount = 1;
- tp.Privileges[0].Attributes = (enabled ? SE_PRIVILEGE_ENABLED : 0);
-
- if (LookupPrivilegeValueW(NULL, privilegeName, &tp.Privileges[0].Luid))
- {
- if (AdjustTokenPrivileges(hToken, FALSE, &tp, 0, NULL, NULL))
- {
- if (GetLastError() == ERROR_NOT_ALL_ASSIGNED)
- {
- error = ERROR_NOT_ALL_ASSIGNED;
- goto fail;
- }
- }
- else
- {
- error = GetLastError();
- goto fail;
- }
- }
- else
- {
- error = GetLastError();
- goto fail;
- }
- }
- else
- {
- error = GetLastError();
- goto fail;
- }
-
- CloseHandle(hToken);
- return TRUE;
-
-fail:
- // Don't allow a success error to be printed, as that would confuse the user.
- // "Access denied" seems like a reasonable default.
- if (error == ERROR_SUCCESS) error = ERROR_ACCESS_DENIED;
- if (hToken != INVALID_HANDLE_VALUE) CloseHandle(hToken);
-
- output_error(error);
- return FALSE;
-}
-
-static int reg_save(int argc, WCHAR* argv[]) {
- HKEY root, hkey;
- LSTATUS status;
- WCHAR* path, *long_key;
-
- if (argc < 4 || argc > 5) goto error;
-
- if (!parse_registry_key(argv[2], &root, &path, &long_key))
- return 1;
-
- if (GetFileAttributes(argv[3]) != INVALID_FILE_ATTRIBUTES)
- {
- if (argc == 5 && !strcmpiW(argv[4], L"/y"))
- {
- DeleteFile(argv[3]);
- }
- else
- {
- if (ask_confirm(STRING_OVERWRITE_FILE, argv[3]))
- DeleteFile(argv[3]);
- }
- }
-
- if (RegOpenKeyExW(root, path, 0, KEY_READ, &hkey))
- {
- output_message(STRING_INVALID_KEY);
- return 1;
- }
-
- if (!set_privilege(SE_BACKUP_NAME, TRUE)) return 1;
-
- status = RegSaveKeyExW(hkey, argv[3], NULL, REG_LATEST_FORMAT);
- RegCloseKey(hkey);
-
- if (status != ERROR_SUCCESS) {
- output_error(status);
- return 1;
- }
-
- return 0;
-
-error:
- output_message(STRING_INVALID_SYNTAX);
- output_message(STRING_FUNC_HELP, struprW(argv[1]));
- return 1;
-}
-
-static int reg_restore(int argc, WCHAR* argv[])
-{
- HKEY root, hkey;
- LSTATUS status;
- WCHAR* path, * long_key;
-
- if (argc != 4) goto error;
-
- if (!parse_registry_key(argv[2], &root, &path, &long_key))
- return 1;
-
- if (RegOpenKeyExW(root, path, 0, KEY_READ, &hkey))
- {
- output_message(STRING_INVALID_KEY);
- return 1;
- }
-
- if (!set_privilege(SE_BACKUP_NAME, TRUE)) return 1;
- if (!set_privilege(SE_RESTORE_NAME, TRUE)) return 1;
-
- status = RegRestoreKeyW(hkey, argv[3], 0);
- RegCloseKey(hkey);
-
- if (status != ERROR_SUCCESS) {
- output_error(status);
- return 1;
- }
-
- return 0;
-
-error:
- output_message(STRING_INVALID_SYNTAX);
- output_message(STRING_FUNC_HELP, struprW(argv[1]));
- return 1;
-}
-