/* -*- Mode: rust; rust-indent-offset: 4 -*- */
/* 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/. */

//! # Limited Access Feature Service
//!
//! Limited Access Features are an interface used by Windows to gate access to
//! APIs. This module implements utilities to unlock features as well as exposes
//! this functionality through XPCOM.
//!
//! ---
//!
//! This module implements the same method used by Microsoft to generate the
//! Limited Access Feature token. The following is provided for additional
//! context.
//!
//! ## Microsoft Token Generation
//!
//! To unlock features, we need:
//! - a feature identifier
//! - a token
//! - an attestation string
//!
//! The token can be 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 documented on MSDN
//! for the feature you want unlocked to contact.
//!
//! 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
//! Windows features such as the taskbar pinning APIs from an unpackaged build.
//!
//! In the above, the token is generated from the Publisher ID (pcsmm0jrprpb2)
//! and the Product Name (MozillaFirefox)
//!
//! ## Microsoft Provided Tokens
//!
//! All tokens listed here were provided to us by Microsoft. Per Microsoft,
//! these tokens are not considered secret, thus have been included in source.
//!
//! Below and in create_rc.py, we used this set:
//!
//! > Product Name: "MozillaFirefox"
//! > Publisher ID: "pcsmm0jrprpb2"
//! > Token: "kRFiWpEK5uS6PMJZKmR7MQ=="
//!
//! Microsoft also provided these other tokens, which will work if accompanied
//! by the matching changes to create_rc.py:
//!
//! > Product Name: "FirefoxBeta"
//! > Publisher ID: "pcsmm0jrprpb2"
//! > Token: "RGEhsYgKhmPLKyzkEHnMhQ=="
//!
//! > Product Name: "FirefoxNightly"
//! > Publisher ID: "pcsmm0jrprpb2"
//! > Token: "qbVzns/9kT+t15YbIwT4Jw=="
//!
//! 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 use the matching Token and Attestation
//! strings.
//!
//! 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 [Microsoft examples](https://github.com/microsoft/Windows-classic-samples/tree/main/Samples/TaskbarManager/CppUnpackagedDesktopTaskbarPin)

use base64::{Engine as _, engine::general_purpose::STANDARD};
use log;
use nserror::{NS_ERROR_UNEXPECTED, NS_OK, nsresult};
use nsstring::{nsACString, nsCString, nsString};
use sha2::{Digest, Sha256};
use std::borrow::Cow;
use windows::{
    ApplicationModel::{LimitedAccessFeatureStatus as LafStatus, LimitedAccessFeatures, Package},
    Win32::Foundation::APPMODEL_ERROR_NO_PACKAGE,
    core::{HRESULT, HSTRING},
};
use xpcom::{
    RefPtr,
    interfaces::{nsILimitedAccessFeature, nsIWindowsRegKey},
    xpcom, xpcom_method,
};

#[xpcom(implement(nsILimitedAccessFeature), nonatomic)]
struct LimitedAccessFeature {
    feature_id: String,
    token: String,
    attestation: String,
}

impl LimitedAccessFeature {
    xpcom_method!(unlock => Unlock() -> bool);
    pub fn unlock(&self) -> Result<bool, nsresult> {
        let status = LimitedAccessFeatures::TryUnlockFeature(
            &HSTRING::from(&self.feature_id),
            &HSTRING::from(&self.token),
            &HSTRING::from(&self.attestation),
        )
        .and_then(|result| result.Status())
        .map_err(|e| {
            log::error!("{e:?}");
            NS_ERROR_UNEXPECTED
        })?;

        log::debug!("Unlock status: {status:?}");

        match status {
            LafStatus::Available | LafStatus::AvailableWithoutToken => Ok(true),
            LafStatus::Unavailable | LafStatus::Unknown => Ok(false),
            _ => Err(NS_ERROR_UNEXPECTED),
        }
    }

    xpcom_method!(get_feature_id => GetFeatureId() -> nsACString);
    fn get_feature_id(&self) -> Result<nsCString, nsresult> {
        Ok(self.feature_id.as_str().into())
    }

    xpcom_method!(get_token => GetToken() -> nsACString);
    fn get_token(&self) -> Result<nsCString, nsresult> {
        Ok(self.token.as_str().into())
    }

    xpcom_method!(get_attestation => GetAttestation() -> nsACString);
    fn get_attestation(&self) -> Result<nsCString, nsresult> {
        Ok(self.attestation.as_str().into())
    }
}

#[xpcom(implement(nsILimitedAccessFeatureService), atomic)]
pub struct LimitedAccessFeatureService {}

impl LimitedAccessFeatureService {
    xpcom_method!(get_taskbar_pin_feature_id => GetTaskbarPinFeatureId() -> nsACString);
    pub fn get_taskbar_pin_feature_id(&self) -> Result<nsCString, nsresult> {
        Ok(nsCString::from("com.microsoft.windows.taskbar.pin"))
    }

    xpcom_method!(generate_limited_access_feature => GenerateLimitedAccessFeature(featureId: *const nsACString) -> *const nsILimitedAccessFeature);
    pub fn generate_limited_access_feature(
        &self,
        feature_id: &nsACString,
    ) -> Result<RefPtr<nsILimitedAccessFeature>, nsresult> {
        let (family_name, publisher_id) = get_package_identity()?;
        let token = generate_token(&feature_id.to_utf8(), &family_name)
            .inspect_err(|e| log::error!("Error generating feature token: {e:?}"))?;
        let attestation = generate_attestation(&feature_id.to_utf8(), &publisher_id);

        let feature = LimitedAccessFeature::allocate(InitLimitedAccessFeature {
            feature_id: feature_id.to_utf8().into(),
            token,
            attestation: attestation,
        });

        feature
            .query_interface::<nsILimitedAccessFeature>()
            .ok_or(NS_ERROR_UNEXPECTED)
    }
}

// Generates a token for a given feature ID.
//
// This function first retrieves the LAF key from the registry using the LAF
// identifier and then combines the feature_id, feature_key, and family_name
// into a token.
//
// Base64(SHA256Encode("{feature_id}!{feature_key}!{family_name}")[0..16])
// yields the complete LAF token for unlocking.
fn generate_token(feature_id: &str, family_name: &FamilyName) -> Result<String, nsresult> {
    let family_name = &family_name.0;
    let feature_key = get_feature_key(feature_id)?;

    let to_hash = format!("{feature_id}!{feature_key}!{family_name}");

    let mut hasher = Sha256::new();
    hasher.update(to_hash);
    let digest = hasher.finalize();
    Ok(STANDARD.encode(&digest[..16]))
}

// Retrieves the feature key from the Windows registry.
fn get_feature_key(feature_id: &str) -> Result<String, nsresult> {
    let reg: RefPtr<nsIWindowsRegKey> =
        xpcom::create_instance(c"@mozilla.org/windows-registry-key;1")
            .ok_or(NS_ERROR_UNEXPECTED)?;

    let path = nsString::from(&format!(
        r"SOFTWARE\Microsoft\Windows\CurrentVersion\AppModel\LimitedAccessFeatures\{feature_id}"
    ));
    // SAFETY: `path` was initialized above.
    unsafe {
        reg.Open(
            nsIWindowsRegKey::ROOT_KEY_LOCAL_MACHINE,
            &*path,
            nsIWindowsRegKey::ACCESS_QUERY_VALUE,
        )
    }
    .to_result()?;

    let mut value = nsString::new();
    // SAFETY: `value` was initialized above, other arguments are initialized inline.
    unsafe { reg.ReadStringValue(&*nsString::new(), &mut *value) }.to_result()?;

    Ok(value.to_string())
}

fn generate_attestation(feature_id: &str, publisher_id: &PublisherId) -> String {
    let publisher_id = &publisher_id.0;
    format!(
        "{publisher_id} has registered their use of {feature_id} with Microsoft and agrees to the terms of use."
    )
}

// The Family Name associated to the application.
//
// For packaged MSIX installs this is derived as `"{Name}_{Publisher ID}"` from
// `<Identity Name="" Publisher="">` in AppxManfiest.xml. Note that `Publisher`
// is not equivalent to `PublisherId`.
//
// For unpackaged applications this is set by `Identity LimitedAccessFeature` in
// the embedded .rc resource file.
struct FamilyName<'a>(Cow<'a, str>);

// The Publisher ID associated to the application.
//
// For packaged MSIX installs this is derived as a hash of `<Identity
// Publisher="">` from AppxManifest.xml.
//
// For unpackaged applications this this is inferred from last 13 characters of
// Package Family Name defined by `Identity LimitedAccessFeature` in the
// embedded .rc resource file.
struct PublisherId<'a>(Cow<'a, str>);

// Retrieves the Package Identity - Family Name and Publisher ID - necessary to
// unlock a limited access feature.
fn get_package_identity<'a>() -> Result<(FamilyName<'a>, PublisherId<'a>), nsresult> {
    match Package::Current() {
        Ok(package) => (|| {
            let id = package.Id()?;
            let family_name = FamilyName(id.FamilyName()?.to_string().into());
            let publisher_id = PublisherId(id.PublisherId()?.to_string().into());
            Ok((family_name, publisher_id))
        })()
        .map_err(|e: HRESULT| {
            log::error!("{e:?}");
            NS_ERROR_UNEXPECTED
        }),
        Err(e) if e.code() == APPMODEL_ERROR_NO_PACKAGE.to_hresult() => {
            // The non-MSIX family name must match that set for `Identity
            // LimitedAccessFeature` in the resource file generated by
            // create_rc.py.
            const UNPACKAGED_FAMILY_NAME: &str = "MozillaFirefox_pcsmm0jrprpb2";
            const UNPACKAGED_PUBLISHER_ID: &str = "pcsmm0jrprpb2";

            log::debug!(
                "Not an MSIX install, using Family Name: `{UNPACKAGED_FAMILY_NAME}` and Publisher ID: `{UNPACKAGED_PUBLISHER_ID}`"
            );
            Ok((
                FamilyName(UNPACKAGED_FAMILY_NAME.into()),
                PublisherId(UNPACKAGED_PUBLISHER_ID.into()),
            ))
        }
        Err(e) => {
            log::error!("{e:?}");
            Err(NS_ERROR_UNEXPECTED)
        }
    }
}

/// Constructor to allow the `nsILimitedAccessFeatureService` to be created through the C ABI.
///
/// # Safety
///
/// This function much be called with valid `iid` and `result` pointers.
#[unsafe(no_mangle)]
pub extern "C" fn new_limited_access_feature_service(
    iid: *const xpcom::nsIID,
    result: *mut *mut xpcom::reexports::libc::c_void,
) -> nsresult {
    let service = LimitedAccessFeatureService::allocate(InitLimitedAccessFeatureService {});
    // SAFETY: The caller is responsible to pass a valid IID and pointer-to-pointer.
    unsafe { service.QueryInterface(iid, result) }
}