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

#include "chrome/browser/ash/arc/policy/managed_configuration_variables.h"

#include <string>
#include <string_view>
#include <vector>

#include "base/check.h"
#include "base/compiler_specific.h"
#include "base/containers/flat_map.h"
#include "base/functional/bind.h"
#include "base/strings/string_split.h"
#include "base/strings/string_util.h"
#include "base/strings/stringprintf.h"
#include "base/values.h"
#include "chrome/browser/ash/policy/core/device_attributes.h"
#include "chromeos/ash/components/system/statistics_provider.h"
#include "components/user_manager/user.h"
#include "third_party/re2/src/re2/re2.h"

namespace arc {

namespace {

// Part before "@" of the given |email| address.
// "some_email@domain.com" => "some_email"
//
// Returns empty string if |email| does not contain an "@".
std::string EmailName(const std::string& email) {
  size_t at_sign_pos = email.find("@");
  if (at_sign_pos == std::string::npos) {
    return "";
  }
  return email.substr(0, at_sign_pos);
}

// Part after "@" of an email address.
// "some_email@domain.com" => "domain.com"
//
// Returns empty string if |email| does not contain an "@".
std::string EmailDomain(const std::string& email) {
  size_t at_sign_pos = email.find("@");
  if (at_sign_pos == std::string::npos) {
    return "";
  }
  return email.substr(at_sign_pos + 1);
}

std::string Identity(const std::string& email) {
  return email;
}

std::string DeviceSerialNumber() {
  return std::string(
      ash::system::StatisticsProvider::GetInstance()->GetMachineID().value_or(
          ""));
}

// Map associating known variables to functions that return the corresponding
// values. For example, "USER_EMAIL" is associated to |SignedInUserEmail|.
typedef base::flat_map<std::string, base::RepeatingCallback<std::string()>>
    VariableResolver;

// Builds a `VariableResolver` from all known variables.
// `attributes` must not be null and must outlive the returned
// `VariableResolver`.
const VariableResolver BuildVariableResolver(
    const user_manager::User& user,
    const policy::DeviceAttributes* attributes) {
  CHECK(attributes);

  // Use |empty_string_getter| for device attributes if user is not affiliated.
  const bool is_affiliated = user.IsAffiliated();
  const std::string user_email = user.GetAccountId().GetUserEmail();
  const auto empty_string_getter =
      base::BindRepeating([]() { return std::string(); });

  return VariableResolver{
      {kUserEmail, base::BindRepeating(&Identity, user_email)},
      {kUserEmailName, base::BindRepeating(&EmailName, user_email)},
      {kUserEmailDomain, base::BindRepeating(&EmailDomain, user_email)},
      {kDeviceDirectoryId,
       is_affiliated
           // The safety of Unretained is upheld by the caller.
           ? base::BindRepeating(&policy::DeviceAttributes::GetDirectoryApiID,
                                 base::Unretained(attributes))
           : empty_string_getter},
      {kDeviceSerialNumber, is_affiliated
                                ? base::BindRepeating(&DeviceSerialNumber)
                                : empty_string_getter},
      {kDeviceAssetId,
       is_affiliated
           // The safety of Unretained is upheld by the caller.
           ? base::BindRepeating(&policy::DeviceAttributes::GetDeviceAssetID,
                                 base::Unretained(attributes))
           : empty_string_getter},
      {kDeviceAnnotatedLocation,
       is_affiliated
           // The safety of Unretained is upheld by the caller.
           ? base::BindRepeating(
                 &policy::DeviceAttributes::GetDeviceAnnotatedLocation,
                 base::Unretained(attributes))
           : empty_string_getter},
  };
}

// Return the value associated to the first item in |variables| that is not
// empty.
std::string ResolveVariableChain(const VariableResolver& resolver,
                                 std::vector<std::string_view> variables) {
  for (const auto& variable : variables) {
    // Variables should always be valid and have a mapping in |resolver|.
    DCHECK(resolver.find(variable) != resolver.end());

    // Resolve the given variable and return if it has a value.
    std::string result = resolver.at(variable).Run();
    if (!result.empty())
      return result;
  }
  return "";
}

std::vector<std::string_view> SplitByColon(std::string_view input) {
  return base::SplitStringPiece(input, ":", base::TRIM_WHITESPACE,
                                base::SPLIT_WANT_NONEMPTY);
}

// Return a new string where all captures of |regex| in |search_input| have been
// replaced with the output of |replacement_getter.Run(capture)|.
std::string SearchAndReplace(
    const re2::RE2& regex,
    base::RepeatingCallback<std::string(std::string_view)> replacement_getter,
    std::string_view search_input) {
  std::vector<std::string> output;
  std::string_view capture;

  // Loop as long as |regex| matches |search_input|.
  while (re2::RE2::PartialMatch(search_input, regex, &capture)) {
    DCHECK(capture.data() != nullptr);
    // Output the prefix skipped by PartialMatch until |capture| is found.
    DCHECK(capture.begin() >= search_input.begin());
    size_t prefix_size = capture.begin() - search_input.begin();
    output.emplace_back(search_input.data(), prefix_size);
    // Output the replacement for |capture|.
    output.emplace_back(replacement_getter.Run(capture));

    // Update |search_input| to the suffix after |capture|.
    DCHECK(search_input.length() >= prefix_size + capture.length());
    size_t remaining_size =
        search_input.length() - (prefix_size + capture.length());
    search_input = std::string_view(
        UNSAFE_TODO(capture.data() + capture.size()), remaining_size);
  }
  // Output the remaining |search_input|.
  output.emplace_back(search_input);
  return base::JoinString(output, /*separator=*/"");
}

// Returns a regular expression that matches any one variable in |resolver|.
std::string ResolverKeyMatcher(const VariableResolver& resolver) {
  std::vector<std::string_view> keys;
  for (const auto& item : resolver) {
    keys.emplace_back(item.first);
  }
  return base::JoinString(keys, /*separator=*/"|");
}

// Replace all variable chains in |configuration| in-place using the provided
// |resolver| rules.
//
// A variable chain is separated by ":", e.g. "${DEVICE_ASSET_ID:USER_EMAIL}".
//
// Chains resolve to the first value that is non-empty. In the example above if
// the asset ID is empty, the chain resolves to the email of the current user.
void ReplaceVariables(const VariableResolver& resolver,
                      std::string& configuration) {
  // |variable_matcher| matches any of the supported variables in |resolver|.
  const std::string variable_matcher = ResolverKeyMatcher(resolver);

  // |variable_capture| will match and capture a variable template including a
  // variable chain. This regex does not match templates with invalid variables.
  const std::string variable_capture =
      base::StringPrintf("(\\$\\{(?:%s)(?::(?:%s))*\\})",
                         variable_matcher.c_str(), variable_matcher.c_str());
  const re2::RE2 regex(variable_capture);
  DCHECK(regex.ok()) << "Error compiling regex: " << regex.error();

  // Callback to compute values of variable chains matched with |regex|.
  auto chain_resolver = base::BindRepeating(
      [](const VariableResolver& resolver, std::string_view variable) {
        // Remove the "${" prefix and the "}" suffix from |variable|.
        DCHECK(variable.starts_with("${") && variable.ends_with("}"));
        const std::string_view chain = variable.substr(2, variable.size() - 3);
        const std::vector<std::string_view> variables = SplitByColon(chain);

        const std::string chain_value =
            ResolveVariableChain(resolver, variables);

        return chain_value;
      },
      resolver);

  std::string replaced_configuration =
      SearchAndReplace(regex, std::move(chain_resolver), configuration);
  configuration = std::move(replaced_configuration);
}

void ReplaceVariables(const VariableResolver& resolver,
                      base::Value& configuration);

void ReplaceVariables(const VariableResolver& resolver,
                      base::DictValue& configuration) {
  for (auto entry : configuration) {
    ReplaceVariables(resolver, entry.second);
  }
}

void ReplaceVariables(const VariableResolver& resolver,
                      base::ListValue& configuration) {
  for (auto& entry : configuration) {
    ReplaceVariables(resolver, entry);
  }
}

void ReplaceVariables(const VariableResolver& resolver,
                      base::Value& configuration) {
  if (configuration.is_dict()) {
    ReplaceVariables(resolver, configuration.GetDict());
  } else if (configuration.is_string()) {
    ReplaceVariables(resolver, configuration.GetString());
  } else if (configuration.is_list()) {
    ReplaceVariables(resolver, configuration.GetList());
  }
}

}  // namespace

const char kUserEmail[] = "USER_EMAIL";
const char kUserEmailName[] = "USER_EMAIL_NAME";
const char kUserEmailDomain[] = "USER_EMAIL_DOMAIN";
const char kDeviceDirectoryId[] = "DEVICE_DIRECTORY_ID";
const char kDeviceSerialNumber[] = "DEVICE_SERIAL_NUMBER";
const char kDeviceAssetId[] = "DEVICE_ASSET_ID";
const char kDeviceAnnotatedLocation[] = "DEVICE_ANNOTATED_LOCATION";

void RecursivelyReplaceManagedConfigurationVariables(
    const user_manager::User& user,
    const policy::DeviceAttributes& device_attributes,
    base::DictValue& managedConfiguration) {
  const VariableResolver resolver =
      BuildVariableResolver(user, &device_attributes);
  ReplaceVariables(resolver, managedConfiguration);
}

}  // namespace arc