// Copyright 2024 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "base/win/elevation_util.h"

#include <objbase.h>

#include <windows.h>

#include <shlobj.h>
#include <wrl/client.h>

#include <string>
#include <utility>

#include "base/base_paths.h"
#include "base/command_line.h"
#include "base/files/file_path.h"
#include "base/logging.h"
#include "base/path_service.h"
#include "base/process/launch.h"
#include "base/process/process.h"
#include "base/process/process_handle.h"
#include "base/process/process_info.h"
#include "base/win/access_token.h"
#include "base/win/scoped_bstr.h"
#include "base/win/scoped_process_information.h"
#include "base/win/scoped_variant.h"
#include "base/win/shell_util.h"
#include "base/win/startup_information.h"
#include "third_party/abseil-cpp/absl/cleanup/cleanup.h"

namespace base::win {

ProcessId GetExplorerPid() {
  const HWND hwnd = ::GetShellWindow();
  ProcessId pid = 0;
  return hwnd && ::GetWindowThreadProcessId(hwnd, &pid) ? pid : kNullProcessId;
}

bool IsProcessRunningAtMediumOrLower(ProcessId process_id) {
  IntegrityLevel level = GetProcessIntegrityLevel(process_id);
  return level != INTEGRITY_UNKNOWN && level <= MEDIUM_INTEGRITY;
}

bool IsProcessRunningSplitToken(ProcessId process_id) {
  auto process =
      Process::OpenWithAccess(process_id, PROCESS_QUERY_LIMITED_INFORMATION);
  if (!process.IsValid()) {
    return false;
  }
  std::optional<win::AccessToken> token =
      AccessToken::FromProcess(process.Handle());
  return token && token->IsSplitToken();
}

expected<Process, DWORD> LaunchProcessDirectly(
    const CommandLine& command_line) {
  LaunchOptions options;
  options.grant_foreground_privilege = true;
  if (auto process = LaunchProcess(command_line, options); process.IsValid()) {
    return ok(std::move(process));
  }
  return unexpected(::GetLastError());
}

// Based on
// https://learn.microsoft.com/en-us/archive/blogs/aaron_margosis/faq-how-do-i-start-a-program-as-the-desktop-user-from-an-elevated-app.
expected<Process, DWORD> RunDeElevated(
    const CommandLine& command_line,
    std::optional<ProcessId> medium_process_id) {
  if (!::IsUserAnAdmin()) {
    return LaunchProcessDirectly(command_line);
  }

  const ProcessId medium_pid =
      medium_process_id ? *medium_process_id : GetExplorerPid();
  if (!medium_pid) {
    return unexpected(static_cast<DWORD>(ERROR_ACCESS_DENIED));
  }
  if (!IsProcessRunningSplitToken(medium_pid)) {
    return LaunchProcessDirectly(command_line);
  }
  if (!IsProcessRunningAtMediumOrLower(medium_pid)) {
    return unexpected(static_cast<DWORD>(ERROR_ACCESS_DENIED));
  }

  auto shell_process =
      Process::OpenWithAccess(medium_pid, PROCESS_QUERY_LIMITED_INFORMATION);
  if (!shell_process.IsValid()) {
    return unexpected(::GetLastError());
  }

  auto token = AccessToken::FromProcess(
      ::GetCurrentProcess(), /*impersonation=*/false, MAXIMUM_ALLOWED);
  if (!token) {
    return unexpected(::GetLastError());
  }
  auto previous_impersonate = token->SetPrivilege(SE_IMPERSONATE_NAME, true);
  if (!previous_impersonate) {
    return unexpected(::GetLastError());
  }
  absl::Cleanup restore_previous_privileges = [&] {
    token->SetPrivilege(SE_IMPERSONATE_NAME, *previous_impersonate);
  };

  auto shell_token = AccessToken::FromProcess(
      shell_process.Handle(), /*impersonation=*/false, TOKEN_DUPLICATE);
  if (!shell_token) {
    return unexpected(::GetLastError());
  }

  auto duplicated_shell_token = shell_token->DuplicatePrimary(
      TOKEN_QUERY | TOKEN_ASSIGN_PRIMARY | TOKEN_DUPLICATE |
      TOKEN_ADJUST_DEFAULT | TOKEN_ADJUST_SESSIONID);
  if (!duplicated_shell_token) {
    return unexpected(::GetLastError());
  }

  StartupInformation startupinfo;
  PROCESS_INFORMATION pi = {};
  if (!::CreateProcessWithTokenW(duplicated_shell_token->get(), 0,
                                 command_line.GetProgram().value().c_str(),
                                 command_line.GetCommandLineString().data(), 0,
                                 nullptr, nullptr, startupinfo.startup_info(),
                                 &pi)) {
    return unexpected(::GetLastError());
  }
  ScopedProcessInformation process_info(pi);
  Process process(process_info.TakeProcessHandle());
  const DWORD pid = process.Pid();
  VLOG(1) << __func__ << ": Started process, PID: " << pid;

  // Allow the spawned process to show windows in the foreground.
  if (!::AllowSetForegroundWindow(pid)) {
    VPLOG(1) << __func__ << ": ::AllowSetForegroundWindow failed";
  }

  return ok(std::move(process));
}

HRESULT RunDeElevatedNoWait(const CommandLine& command_line) {
  return RunShellExecuteViaExplorer(command_line.GetProgram().value(),
                                    command_line.GetArgumentsString());
}

HRESULT RunDeElevatedNoWait(const std::wstring& path,
                            const std::wstring& parameters,
                            std::optional<std::wstring_view> current_directory,
                            bool start_hidden) {
  ShellExecuteOptions options{.current_directory = std::wstring(
                                  current_directory.value_or(std::wstring())),
                              .start_hidden = start_hidden};
  return RunShellExecuteViaExplorer(path, parameters, options);
}

}  // namespace base::win