/* -*- Mode: C++; tab-width: 2; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

/**
 * This file exists so that LaunchModernSettingsDialogDefaultApps can be called
 * without linking to libxul.
 */
#include "Windows11LimitedAccessFeatures.h"

#include "mozilla/Logging.h"

static mozilla::LazyLogModule sLog("Windows11LimitedAccessFeatures");

#define LAF_LOG(level, msg, ...) MOZ_LOG(sLog, level, (msg, ##__VA_ARGS__))

// MINGW32 is not supported for these features
// Fall back function defined in the #else
#ifndef __MINGW32__

#  include "mozilla/ErrorResult.h"
#  include "nsString.h"
#  include "nsCOMPtr.h"
#  include "nsComponentManagerUtils.h"
#  include "nsIWindowsRegKey.h"
#  include "nsWindowsHelpers.h"
#  include "nsICryptoHash.h"

#  include "mozilla/Atomics.h"
#  include "mozilla/Base64.h"
#  include "mozilla/Char16.h"
#  include "mozilla/WinHeaderOnlyUtils.h"
#  include "WinUtils.h"

#  include <wrl.h>
#  include <inspectable.h>
#  include <roapi.h>
#  include <windows.services.store.h>
#  include <windows.foundation.h>

using namespace Microsoft::WRL;
using namespace Microsoft::WRL::Wrappers;
using namespace ABI::Windows;
using namespace ABI::Windows::Foundation;
using namespace ABI::Windows::ApplicationModel;

using namespace mozilla;

/**
 * To unlock features, we need:
 * a feature identifier
 * a token,
 * an attestation string
 * a token
 *
 * The token is generated by Microsoft and must
 * match the publisher id Microsoft thinks we have, for a particular
 * feature.
 *
 * To get a token, find the right microsoft email address by doing
 * a search on the web for the feature you want unlocked and reach
 * out to the right people at Microsoft.
 *
 * The token is generated from Microsoft.
 * The jumbled code in the attestation string is a publisher id and
 * must match the code in the resources / .rc file for the identity,
 * looking like this for non-MSIX builds:
 *
 * Identity LimitedAccessFeature {{ L"MozillaFirefox_pcsmm0jrprpb2" }}
 *
 * Broken down:
 * Identity LimitedAccessFeature {{ L"PRODUCTNAME_PUBLISHERID" }}
 *
 * That is injected into our build in create_rc.py and is necessary
 * to unlock the taskbar pinning feature / APIs from an unpackaged
 * build.
 *
 * In the above, the token is generated from the publisher id (pcsmm0jrprpb2)
 * and the product name (MozillaFirefox)
 *
 * All tokens listed here were provided to us by Microsoft.
 *
 * Below and in create_rc.py, we used this set:
 *
 * Token: "kRFiWpEK5uS6PMJZKmR7MQ=="
 * Product Name: "MozillaFirefox"
 * Publisher ID: "pcsmm0jrprpb2"
 *
 * Microsoft also provided these other tokens, which will will
 * work if accompanied by the matching changes to create_rc.py:

 * -----
 * Token: "RGEhsYgKhmPLKyzkEHnMhQ=="
 * Product Name: "FirefoxBeta"
 * Publisher ID: "pcsmm0jrprpb2"
 *
 * -----
 *
 * Token: "qbVzns/9kT+t15YbIwT4Jw=="
 * Product Name: "FirefoxNightly"
 * Publisher ID: "pcsmm0jrprpb2"
 *
 * To use those instead, you have to ensure that the LimitedAccessFeature
 * generated in create_rc.py has the product name and publisher id
 * matching the token used in this file.
 *
 * For non-packaged (non-MSIX) builds, any of the above sets will work.
 * Just make sure the right (ProductName_PublisherID) value is in the
 * generated resource data for the executable, and the matching
* (Token) and attestation string
 *
 * To get MSIX/packaged builds to work, the product name and publisher in
 * the final manifest (searchfox.org/mozilla-central/search?q=APPX_PUBLISHER)
 * should match the token in this file. For that case, the identity value
 * in the resources does not matter.
 *
 * See here for Microsoft examples:
https://github.com/microsoft/Windows-classic-samples/tree/main/Samples/TaskbarManager/CppUnpackagedDesktopTaskbarPin
 */

/**
 * Unlocks a Windows Limited Access Feature (LAF) by generating a token and
 * attestation.
 *
 * This function first retrieves the LAF key from the registry using
 * the LAF identifier and then combines the lafId, lafKey, and PFN
 * into a token.
 *
 * Applying Base64(SHA256Encode("<lafId>!<lafKey>!<PFN>")[0..16]) yields the
 * complete LAF token for unlocking.
 *
 * Taking the last 13 characters of the PFN yields the publisher identifier
 * which is used in the following boilerplate:
 * "<PFN[-13]> has registered their use of <lafId> with Microsoft and
 * agrees to the terms of use."
 *
 * @return {LimitedAccessFeatureInfo} containing the generated
 *   token and attestation upon success. Contains empty strings for
 *   these fields upon failure.
 */
static mozilla::Result<LimitedAccessFeatureInfo, nsresult>
GenerateLimitedAccessFeatureInfo(const nsCString& debugName,
                                 const nsString& lafId) {
  nsresult rv;
  // Read registry key for a given Limited Access Feature with ID lafId.
  nsAutoString keyData;
  nsCOMPtr<nsIWindowsRegKey> regKey =
      do_CreateInstance("@mozilla.org/windows-registry-key;1", &rv);
  NS_ENSURE_SUCCESS(rv, Err(rv));
  const nsAutoString regPath =
      u"SOFTWARE\\Microsoft\\Windows\\CurrentVersion\\AppModel\\LimitedAccessFeatures\\"_ns +
      lafId;
  rv = regKey->Open(nsIWindowsRegKey::ROOT_KEY_LOCAL_MACHINE, regPath,
                    nsIWindowsRegKey::ACCESS_READ);
  NS_ENSURE_SUCCESS(rv, Err(rv));
  rv = regKey->ReadStringValue(u""_ns, keyData);
  NS_ENSURE_SUCCESS(rv, Err(rv));

  // Get Package Family Name (PFN) and assemble a string to hash.
  // First convert this string to SHA256, take the first 16 bytes
  // only, and then read as base 64.
  nsAutoCString hashStringResult;
  nsAutoString encodedToken;
  // The non-MSIX family name must match whatever value is in create_rc.py
  // Currently this is MozillaFirefox_pcsmm0jrprpb2
  nsAutoString familyName;
  if (widget::WinUtils::HasPackageIdentity()) {
    familyName = nsDependentString(mozilla::GetPackageFamilyName().get());
  } else {
    familyName = u"MozillaFirefox_pcsmm0jrprpb2"_ns;
  }
  const nsAutoCString hashString =
      NS_ConvertUTF16toUTF8(lafId + u"!"_ns + keyData + u"!"_ns + familyName);

  nsCOMPtr<nsICryptoHash> cryptoHash =
      do_CreateInstance("@mozilla.org/security/hash;1", &rv);
  NS_ENSURE_SUCCESS(rv, Err(rv));
  rv = cryptoHash->Init(nsICryptoHash::SHA256);
  NS_ENSURE_SUCCESS(rv, Err(rv));
  rv = cryptoHash->Update(reinterpret_cast<const uint8_t*>(hashString.get()),
                          hashString.Length());
  NS_ENSURE_SUCCESS(rv, Err(rv));
  rv = cryptoHash->Finish(false, hashStringResult);
  NS_ENSURE_SUCCESS(rv, Err(rv));

  // Keep only first 16 bytes and encode
  hashStringResult.Truncate(hashStringResult.Length() - 16);
  rv = Base64Encode(hashStringResult, encodedToken);
  NS_ENSURE_SUCCESS(rv, Err(rv));

  // The PFN contains a package ID in the last 13 characters.
  // This ID is based on the value in the publisher field of the
  // AppManifest. This ID is used to assemble the attestation.
  familyName.Cut(0, familyName.Length() - 13);
  nsAutoString attestation =
      familyName + u" has registered their use of "_ns + lafId +
      u" with Microsoft and agrees to the terms of use."_ns;
  LimitedAccessFeatureInfo result = {debugName, lafId, encodedToken,
                                     attestation};
  return result;
}

/**
 Implementation of the Win11LimitedAccessFeaturesInterface.
 */
class Win11LimitedAccessFeatures : public Win11LimitedAccessFeaturesInterface {
 public:
  using AtomicState = Atomic<int, SequentiallyConsistent>;

  Result<bool, HRESULT> Unlock(Win11LimitedAccessFeatureType feature) override;

 private:
  AtomicState& GetState(Win11LimitedAccessFeatureType feature);
  Result<bool, HRESULT> UnlockImplementation(
      const LimitedAccessFeatureInfo& lafInfo);

  /**
   * Store the state as an atomic so that it can be safely accessed from
   * different threads.
   */
  static AtomicState mTaskbarState;
  static AtomicState mDefaultState;

  enum State {
    Uninitialized,
    Locked,
    Unlocked,
  };
};

Win11LimitedAccessFeatures::AtomicState
    Win11LimitedAccessFeatures::mTaskbarState(
        Win11LimitedAccessFeatures::Uninitialized);
Win11LimitedAccessFeatures::AtomicState
    Win11LimitedAccessFeatures::mDefaultState(
        Win11LimitedAccessFeatures::Uninitialized);

RefPtr<Win11LimitedAccessFeaturesInterface>
CreateWin11LimitedAccessFeaturesInterface() {
  RefPtr<Win11LimitedAccessFeaturesInterface> result(
      new Win11LimitedAccessFeatures());
  return result;
}

Result<bool, HRESULT> Win11LimitedAccessFeatures::Unlock(
    Win11LimitedAccessFeatureType feature) {
  AtomicState& atomicState = GetState(feature);

  // Win11LimitedAccessFeatureType::Taskbar
  auto taskbarLafInfo = GenerateLimitedAccessFeatureInfo(
      "Win11LimitedAccessFeatureType::Taskbar"_ns,
      u"com.microsoft.windows.taskbar.pin"_ns);
  if (taskbarLafInfo.isErr()) {
    LAF_LOG(LogLevel::Debug, "Unlocking taskbar failed with error %d",
            NS_ERROR_GET_CODE(taskbarLafInfo.unwrapErr()));
    return Err(E_FAIL);
  }

  LimitedAccessFeatureInfo limitedAccessFeatureInfo[] = {
      taskbarLafInfo.unwrap()};
  const auto& lafInfo = limitedAccessFeatureInfo[static_cast<int>(feature)];

  LAF_LOG(LogLevel::Debug,
          "Limited Access Feature Info for %s. Feature %ls, %ls, %ls",
          lafInfo.debugName.get(), lafInfo.token.getW(), lafInfo.token.getW(),
          lafInfo.attestation.getW());

  int state = atomicState;
  if (state != Uninitialized) {
    LAF_LOG(LogLevel::Debug, "%s already initialized! State = %s",
            lafInfo.debugName.get(), (state == Unlocked) ? "true" : "false");
    return (state == Unlocked);
  }

  // If multiple threads read the state at the same time, and it's unitialized,
  // both threads will unlock the feature. This situation is unlikely, but even
  // if it happens, it's not a problem.

  auto result = UnlockImplementation(lafInfo);

  int newState = Locked;
  if (!result.isErr() && result.unwrap()) {
    newState = Unlocked;
  }

  atomicState = newState;

  return result;
}

Win11LimitedAccessFeatures::AtomicState& Win11LimitedAccessFeatures::GetState(
    Win11LimitedAccessFeatureType feature) {
  switch (feature) {
    case Win11LimitedAccessFeatureType::Taskbar:
      return mTaskbarState;

    default:
      LAF_LOG(LogLevel::Debug, "Missing feature type for %d",
              static_cast<int>(feature));
      MOZ_ASSERT(false,
                 "Unhandled feature type! Add a new atomic state variable, add "
                 "that entry to the switch statement above, and add the proper "
                 "entries for the feature and the token.");
      return mDefaultState;
  }
}

Result<bool, HRESULT> Win11LimitedAccessFeatures::UnlockImplementation(
    const LimitedAccessFeatureInfo& lafInfo) {
  ComPtr<ILimitedAccessFeaturesStatics> limitedAccessFeatures;
  ComPtr<ILimitedAccessFeatureRequestResult> limitedAccessFeaturesResult;

  HRESULT hr = RoGetActivationFactory(
      HStringReference(
          RuntimeClass_Windows_ApplicationModel_LimitedAccessFeatures)
          .Get(),
      IID_ILimitedAccessFeaturesStatics, &limitedAccessFeatures);

  if (!SUCCEEDED(hr)) {
    LAF_LOG(LogLevel::Debug, "%s activation error. HRESULT = 0x%lx",
            lafInfo.debugName.get(), hr);
    return Err(hr);
  }

  hr = limitedAccessFeatures->TryUnlockFeature(
      HStringReference(lafInfo.feature.get()).Get(),
      HStringReference(lafInfo.token.get()).Get(),
      HStringReference(lafInfo.attestation.get()).Get(),
      &limitedAccessFeaturesResult);
  if (!SUCCEEDED(hr)) {
    LAF_LOG(LogLevel::Debug, "%s unlock error. HRESULT = 0x%lx",
            lafInfo.debugName.get(), hr);
    return Err(hr);
  }

  LimitedAccessFeatureStatus status;
  hr = limitedAccessFeaturesResult->get_Status(&status);
  if (!SUCCEEDED(hr)) {
    LAF_LOG(LogLevel::Debug, "%s get status error. HRESULT = 0x%lx",
            lafInfo.debugName.get(), hr);
    return Err(hr);
  }

  int state = Unlocked;
  if ((status != LimitedAccessFeatureStatus_Available) &&
      (status != LimitedAccessFeatureStatus_AvailableWithoutToken)) {
    LAF_LOG(LogLevel::Debug, "%s not available. HRESULT = 0x%lx",
            lafInfo.debugName.get(), hr);
    state = Locked;
  }

  return (state == Unlocked);
}

#else  // MINGW32 implementation

RefPtr<Win11LimitedAccessFeaturesInterface>
CreateWin11LimitedAccessFeaturesInterface() {
  RefPtr<Win11LimitedAccessFeaturesInterface> result;
  return result;
}

static mozilla::Result<LimitedAccessFeatureInfo, nsresult>
GenerateLimitedAccessFeatureInfo(const nsCString& debugName,
                                 const nsString& lafId) {
  return mozilla::Err(NS_ERROR_NOT_IMPLEMENTED);
}

#endif