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

#include <inttypes.h>

#include <iostream>
#include <map>
#include <optional>
#include <set>
#include <string>
#include <string_view>

#include "base/at_exit.h"
#include "base/base_paths.h"
#include "base/command_line.h"
#include "base/compiler_specific.h"
#include "base/containers/span.h"
#include "base/files/file_path.h"
#include "base/files/file_util.h"
#include "base/logging.h"
#include "base/logging/logging_settings.h"
#include "base/numerics/safe_conversions.h"
#include "base/path_service.h"
#include "base/strings/strcat.h"
#include "base/strings/string_number_conversions.h"
#include "base/strings/string_util.h"
#include "base/strings/string_view_util.h"
#include "base/strings/stringprintf.h"
#include "base/strings/utf_string_conversions.h"
#include "build/build_config.h"
#include "crypto/hash.h"
#include "net/cert/root_store_proto_full/root_store.pb.h"
#include "third_party/boringssl/src/include/openssl/bio.h"
#include "third_party/boringssl/src/include/openssl/err.h"
#include "third_party/boringssl/src/include/openssl/pem.h"
#include "third_party/protobuf/src/google/protobuf/text_format.h"

using chrome_root_store::RootStore;
using chrome_root_store::TrustAnchor;

namespace {

// Returns a map from hex-encoded SHA-256 hash to DER certificate, or
// `std::nullopt` if not found.
std::optional<std::map<std::string, std::string>> DecodeCerts(
    std::string_view in) {
  // TODO(crbug.com/40770548): net/cert/pem.h has a much nicer API, but
  // it would require some build refactoring to avoid a circular dependency.
  // This is assuming that the chrome trust store code goes in
  // net/cert/internal, which it may not.
  bssl::UniquePtr<BIO> bio(BIO_new_mem_buf(in.data(), in.size()));
  if (!bio) {
    return std::nullopt;
  }
  std::map<std::string, std::string> certs;
  for (;;) {
    char* name;
    char* header;
    unsigned char* data;
    long len;
    if (!PEM_read_bio(bio.get(), &name, &header, &data, &len)) {
      uint32_t err = ERR_get_error();
      if (ERR_GET_LIB(err) == ERR_LIB_PEM &&
          ERR_GET_REASON(err) == PEM_R_NO_START_LINE) {
        // Found the last PEM block.
        break;
      }
      LOG(ERROR) << "Error reading PEM.";
      return std::nullopt;
    }
    bssl::UniquePtr<char> scoped_name(name);
    bssl::UniquePtr<char> scoped_header(header);
    bssl::UniquePtr<unsigned char> scoped_data(data);
    if (std::string_view(name) != "CERTIFICATE") {
      LOG(ERROR) << "Found PEM block of type " << name
                 << " instead of CERTIFICATE";
      return std::nullopt;
    }
    std::string_view cert_view(reinterpret_cast<const char*>(data),
                               base::checked_cast<size_t>(len));
    const std::string sha256_hex =
        base::HexEncodeLower(crypto::hash::Sha256(cert_view));
    certs[sha256_hex] = std::string(cert_view);
  }
  return std::move(certs);
}

// ReplaceTrustAnchors takes a repeated trust_anchors field from a proto and
// replaces certs identified by a SHA-256 hash with the full cert bytes, using
// the file found at certs_file_path to provide the full cert bytes.
bool ReplaceTrustAnchors(
    google::protobuf::RepeatedPtrField<TrustAnchor>* trust_anchors,
    const base::FilePath& certs_file_path) {
  std::map<std::string, std::string> certs;
  if (!certs_file_path.empty()) {
    std::string certs_data;
    if (!base::ReadFileToString(base::MakeAbsoluteFilePath(certs_file_path),
                                &certs_data)) {
      LOG(ERROR) << "Could not read " << certs_file_path;
      return false;
    }
    auto certs_opt = DecodeCerts(certs_data);
    if (!certs_opt) {
      LOG(ERROR) << "Could not decode " << certs_file_path;
      return false;
    }
    certs = std::move(*certs_opt);
  }

  // Replace the filenames with the actual certificate contents.
  for (auto& anchor : *trust_anchors) {
    if (anchor.certificate_case() !=
        chrome_root_store::TrustAnchor::kSha256Hex) {
      continue;
    }

    auto iter = certs.find(anchor.sha256_hex());
    if (iter == certs.end()) {
      LOG(ERROR) << "Could not find certificate " << anchor.sha256_hex();
      return false;
    }

    // Remove the certificate from `certs`. This both checks for duplicate
    // certificates and allows us to check for unused certificates later.
    anchor.set_der(std::move(iter->second));
    certs.erase(iter);
  }

  if (!certs.empty()) {
    LOG(ERROR) << "Unused certificate (SHA-256 hash " << certs.begin()->first
               << ") in " << certs_file_path;
    return false;
  }
  return true;
}

std::optional<RootStore> ReadTextRootStore(
    const base::FilePath& root_store_path,
    const base::FilePath& certs_path,
    const base::FilePath& additional_certs_path) {
  std::string root_store_text;
  if (!base::ReadFileToString(base::MakeAbsoluteFilePath(root_store_path),
                              &root_store_text)) {
    LOG(ERROR) << "Could not read " << root_store_path;
    return std::nullopt;
  }

  RootStore root_store;
  if (!google::protobuf::TextFormat::ParseFromString(root_store_text,
                                                     &root_store)) {
    LOG(ERROR) << "Could not parse " << root_store_path;
    return std::nullopt;
  }

  if (!ReplaceTrustAnchors(root_store.mutable_trust_anchors(), certs_path)) {
    LOG(ERROR) << "Failed to process " << certs_path;
    return std::nullopt;
  }
  if (!ReplaceTrustAnchors(root_store.mutable_additional_certs(),
                           additional_certs_path)) {
    LOG(ERROR) << "Failed to process " << additional_certs_path;
    return std::nullopt;
  }
  return std::move(root_store);
}

std::string SecondsFromEpochToBaseTime(int64_t t) {
  return base::StrCat({"base::Time::UnixEpoch() + base::Seconds(",
                       base::NumberToString(t), ")"});
}

std::string VersionFromString(std::string_view version_str) {
  return base::StrCat({"\"", version_str, "\""});
}

void WriteByteArrayConstant(std::string_view name,
                            std::string_view data,
                            std::string* string_to_write) {
  base::StringAppendF(string_to_write, "constexpr uint8_t %s[] = {", name);

  // Convert each character to hex representation, escaped.
  for (auto c : data) {
    base::StringAppendF(string_to_write, "0x%02xu,", static_cast<uint8_t>(c));
  }

  *string_to_write += "};\n";
}

void MaybeWriteConstraintsConstant(
    std::string_view name,
    const ::google::protobuf::RepeatedPtrField<
        ::chrome_root_store::ConstraintSet>& constraints,
    std::string* string_to_write) {
  const std::string kNulloptString = "std::nullopt";

  if (constraints.size() <= 0) {
    return;
  }

  int constraint_num = 0;
  for (const auto& constraint : constraints) {
    if (constraint.permitted_dns_names_size() > 0) {
      base::StringAppendF(string_to_write,
                          "constexpr std::string_view "
                          "%sNames%d[] = {",
                          name, constraint_num);
      for (const auto& dns_name : constraint.permitted_dns_names()) {
        base::StringAppendF(string_to_write, "\"%s\",", dns_name);
      }
      *string_to_write += "};\n";
    }
    constraint_num++;
  }

  base::StringAppendF(string_to_write,
                      "constexpr StaticChromeRootCertConstraints "
                      "%s[] = {",
                      name);

  std::vector<std::string> constraint_strings;
  constraint_num = 0;
  for (const auto& constraint : constraints) {
    std::vector<std::string> constraint_params;

    constraint_params.push_back(
        constraint.has_sct_not_after_sec()
            ? SecondsFromEpochToBaseTime(constraint.sct_not_after_sec())
            : kNulloptString);

    constraint_params.push_back(
        constraint.has_sct_all_after_sec()
            ? SecondsFromEpochToBaseTime(constraint.sct_all_after_sec())
            : kNulloptString);

    constraint_params.push_back(
        constraint.has_min_version()
            ? VersionFromString(constraint.min_version())
            : kNulloptString);

    constraint_params.push_back(
        constraint.has_max_version_exclusive()
            ? VersionFromString(constraint.max_version_exclusive())
            : kNulloptString);

    if (constraint.permitted_dns_names_size() > 0) {
      constraint_params.push_back(
          base::StringPrintf("%sNames%d", name, constraint_num));
    } else {
      constraint_params.push_back("{}");
    }

    constraint_params.push_back(
        constraint.has_index_not_after()
            ? base::NumberToString(constraint.index_not_after())
            : kNulloptString);

    constraint_params.push_back(
        constraint.has_index_after()
            ? base::NumberToString(constraint.index_after())
            : kNulloptString);

    constraint_params.push_back(
        constraint.has_validity_starts_not_after_sec()
            ? SecondsFromEpochToBaseTime(
                  constraint.validity_starts_not_after_sec())
            : kNulloptString);

    constraint_params.push_back(
        constraint.has_validity_starts_after_sec()
            ? SecondsFromEpochToBaseTime(constraint.validity_starts_after_sec())
            : kNulloptString);

    constraint_strings.push_back(
        base::StrCat({"{", base::JoinString(constraint_params, ","), "}"}));

    constraint_num++;
  }

  *string_to_write += base::JoinString(constraint_strings, ",");
  *string_to_write += "};\n";
}

void WriteTrustAnchors(
    const google::protobuf::RepeatedPtrField<TrustAnchor>& trust_anchors,
    const std::string& cert_name_prefix,
    std::string* string_to_write) {
  for (int i = 0; i < trust_anchors.size(); i++) {
    const auto& anchor = trust_anchors.Get(i);
    // Every trust anchor at this point should have a DER.
    CHECK(!anchor.der().empty());
    std::string der = anchor.der();

    WriteByteArrayConstant(base::StringPrintf("k%sCert%d", cert_name_prefix, i),
                           der, string_to_write);

    if (!anchor.trust_anchor_id().empty()) {
      WriteByteArrayConstant(
          base::StringPrintf("k%sTrustAnchorID%d", cert_name_prefix, i),
          anchor.trust_anchor_id(), string_to_write);
    }

    MaybeWriteConstraintsConstant(
        base::StringPrintf("k%sConstraints%d", cert_name_prefix, i),
        anchor.constraints(), string_to_write);
  }
}

// Returns true if file was correctly written, false otherwise.
bool WriteRootCppFile(const RootStore& root_store,
                      const base::FilePath cpp_path) {
  // Root store should have at least one trust anchor.
  CHECK_GT(root_store.trust_anchors_size(), 0);

  std::string string_to_write =
      "// This file is auto-generated, DO NOT EDIT.\n\n";

  WriteTrustAnchors(root_store.trust_anchors(), "ChromeRoot", &string_to_write);
  WriteTrustAnchors(root_store.additional_certs(), "Additional",
                    &string_to_write);

  // Assemble list of trust anchors.
  string_to_write += "constexpr ChromeRootCertInfo kChromeRootCertList[] = {\n";
  for (int i = 0; i < root_store.trust_anchors_size(); i++) {
    const auto& anchor = root_store.trust_anchors(i);
    if (anchor.has_tls_trust_anchor()) {
      // Anchors in |trust_anchors| are not supposed to have |tls_trust_anchor|
      // set, because they are definitionally TLS trust anchors.
      return false;
    }
    base::StringAppendF(&string_to_write, "    {kChromeRootCert%d, ", i);
    if (anchor.constraints_size() > 0) {
      base::StringAppendF(&string_to_write, "kChromeRootConstraints%d", i);
    } else {
      string_to_write += "{}";
    }

    base::StringAppendF(
        &string_to_write,
        ", /*enforce_anchor_expiry=*/%s, /*enforce_anchor_constraints=*/%s",
        anchor.enforce_anchor_expiry() ? "true" : "false",
        anchor.enforce_anchor_constraints() ? "true" : "false");

    if (anchor.trust_anchor_id().empty()) {
      string_to_write += ", {}";
    } else {
      base::StringAppendF(&string_to_write, ", kChromeRootTrustAnchorID%d", i);
    }

    if (anchor.has_crs_root_id()) {
      base::StringAppendF(&string_to_write, ", %d", anchor.crs_root_id());
    } else {
      string_to_write += ", std::nullopt";
    }
    string_to_write += "},\n";
  }
  // Append additional_certs as TLS trust anchors, if they are marked as such.
  for (int i = 0; i < root_store.additional_certs_size(); i++) {
    const auto& anchor = root_store.additional_certs(i);
    if (!anchor.tls_trust_anchor()) {
      continue;
    }
    base::StringAppendF(&string_to_write, "    {kAdditionalCert%d, ", i);
    if (anchor.constraints_size() > 0) {
      base::StringAppendF(&string_to_write, "kAdditionalConstraints%d", i);
    } else {
      string_to_write += "{}";
    }

    base::StringAppendF(
        &string_to_write,
        ", /*enforce_anchor_expiry=*/%s, /*enforce_anchor_constraints=*/%s",
        anchor.enforce_anchor_expiry() ? "true" : "false",
        anchor.enforce_anchor_constraints() ? "true" : "false");

    if (anchor.trust_anchor_id().empty()) {
      string_to_write += ", {}";
    } else {
      base::StringAppendF(&string_to_write, ", kAdditionalTrustAnchorID%d", i);
    }

    if (anchor.has_crs_root_id()) {
      base::StringAppendF(&string_to_write, ", %d", anchor.crs_root_id());
    } else {
      string_to_write += ", std::nullopt";
    }
    string_to_write += "},\n";
  }
  string_to_write += "};\n\n";

  // Assemble list of QWAC issuers, which can come from both trust_anchors and
  // additional_certs.
  string_to_write +=
      "constexpr base::span<const uint8_t> kEutlRootCertList[] = {\n";
  for (int i = 0; i < root_store.trust_anchors_size(); i++) {
    const auto& anchor = root_store.trust_anchors(i);
    if (!anchor.eutl()) {
      continue;
    }
    base::StringAppendF(&string_to_write, "    kChromeRootCert%d,\n", i);
  }
  for (int i = 0; i < root_store.additional_certs_size(); i++) {
    const auto& anchor = root_store.additional_certs(i);
    if (!anchor.eutl()) {
      continue;
    }
    base::StringAppendF(&string_to_write, "    kAdditionalCert%d,\n", i);
  }
  string_to_write += "};\n\n";

  // Write constants used by MTC anchors.
  for (int i = 0; i < root_store.mtc_anchors_size(); i++) {
    const auto& anchor = root_store.mtc_anchors(i);
    if (!anchor.tls_trust_anchor()) {
      continue;
    }
    WriteByteArrayConstant(base::StringPrintf("kMtcAnchorLogId%d", i),
                           anchor.log_id(), &string_to_write);
    MaybeWriteConstraintsConstant(
        base::StringPrintf("kMtcAnchorConstraints%d", i), anchor.constraints(),
        &string_to_write);
  }

  // Assemble list of trusted MTC anchors.
  string_to_write +=
      "constexpr ChromeMtcAnchorInfo kChromeTrustedMtcAnchorList[] = {\n";
  for (int i = 0; i < root_store.mtc_anchors_size(); i++) {
    const auto& anchor = root_store.mtc_anchors(i);
    if (!anchor.tls_trust_anchor()) {
      continue;
    }
    base::StringAppendF(&string_to_write, "  {kMtcAnchorLogId%d", i);
    if (anchor.constraints().size() > 0) {
      base::StringAppendF(&string_to_write, ", kMtcAnchorConstraints%d", i);
    } else {
      string_to_write += ", {}";
    }
    if (anchor.has_crs_root_id()) {
      base::StringAppendF(&string_to_write, ", %d", anchor.crs_root_id());
    } else {
      string_to_write += ", std::nullopt";
    }
    string_to_write += "},\n";
  }
  string_to_write += "};\n\n";

  base::StringAppendF(&string_to_write,
                      "\nstatic const int64_t kRootStoreVersion = %" PRId64
                      ";\n",
                      root_store.version_major());
  if (!base::WriteFile(cpp_path, string_to_write)) {
    return false;
  }
  return true;
}

// Returns true if file was correctly written, false otherwise.
bool WriteEvCppFile(const RootStore& root_store,
                    const base::FilePath cpp_path) {
  // There should be at least one EV root.
  CHECK_GT(root_store.trust_anchors_size(), 0);

  std::string string_to_write =
      "// This file is auto-generated, DO NOT EDIT.\n\n"
      "static const EVMetadata kEvRootCaMetadata[] = {\n";

  for (auto& anchor : root_store.trust_anchors()) {
    // Every trust anchor at this point should have a DER.
    CHECK(!anchor.der().empty());
    if (anchor.ev_policy_oids_size() == 0) {
      // The same input file is used for the Chrome Root Store and EV enabled
      // certificates. Skip anchors that have no EV policy OIDs when generating
      // the EV include file.
      continue;
    }

    std::string sha256_hash =
        std::string(base::as_string_view(crypto::hash::Sha256(anchor.der())));

    // Begin struct. Assumed type of EVMetadata:
    //
    // struct EVMetadata {
    //  static const size_t kMaxOIDsPerCA = 2;
    //  SHA256HashValue fingerprint;
    //  const std::string_view policy_oids[kMaxOIDsPerCA];
    // };
    string_to_write += "    {\n";
    string_to_write += "        {{";

    int wrap_count = 0;
    for (auto c : sha256_hash) {
      if (wrap_count != 0) {
        if (wrap_count % 11 == 0) {
          string_to_write += ",\n          ";
        } else {
          string_to_write += ", ";
        }
      }
      base::StringAppendF(&string_to_write, "0x%02x", static_cast<uint8_t>(c));
      wrap_count++;
    }

    string_to_write += "}},\n";
    string_to_write += "        {\n";

    // struct expects exactly two policy oids, and we can only support 1 or 2
    // policy OIDs. These checks will need to change if we ever merge the EV and
    // Chrome Root Store textprotos.
    const int kMaxPolicyOids = 2;
    int oids_size = anchor.ev_policy_oids_size();
    std::string hexencode_hash = base::HexEncode(sha256_hash);
    if (oids_size > kMaxPolicyOids) {
      PLOG(ERROR) << hexencode_hash << " has too many OIDs!";
      return false;
    }
    for (int i = 0; i < kMaxPolicyOids; i++) {
      std::string oid;
      if (i < oids_size) {
        oid = anchor.ev_policy_oids(i);
      }
      base::StrAppend(&string_to_write, {"            \"", oid, "\",\n"});
    }

    // End struct
    string_to_write += "        },\n";
    string_to_write += "    },\n";
  }
  string_to_write += "};\n";
  if (!base::WriteFile(cpp_path, string_to_write)) {
    PLOG(ERROR) << "Error writing cpp include file";
    return false;
  }
  return true;
}

}  // namespace

int main(int argc, char** argv) {
  base::AtExitManager at_exit_manager;
  base::CommandLine::Init(argc, argv);

  logging::LoggingSettings settings;
  settings.logging_dest =
      logging::LOG_TO_SYSTEM_DEBUG_LOG | logging::LOG_TO_STDERR;
  logging::InitLogging(settings);

  base::CommandLine& command_line = *base::CommandLine::ForCurrentProcess();
  base::FilePath proto_path = command_line.GetSwitchValuePath("write-proto");
  base::FilePath root_store_cpp_path =
      command_line.GetSwitchValuePath("write-cpp-root-store");
  base::FilePath ev_roots_cpp_path =
      command_line.GetSwitchValuePath("write-cpp-ev-roots");
  base::FilePath root_store_path =
      command_line.GetSwitchValuePath("root-store");
  base::FilePath certs_path = command_line.GetSwitchValuePath("certs");
  base::FilePath additional_certs_path =
      command_line.GetSwitchValuePath("additional-certs");

  if ((proto_path.empty() && root_store_cpp_path.empty() &&
       ev_roots_cpp_path.empty()) ||
      root_store_path.empty() || command_line.HasSwitch("help")) {
    std::cerr << "Usage: root_store_tool "
              << "--root-store=TEXTPROTO_FILE "
              << "[--certs=CERTS_FILE] "
              << "[--additional-certs=ADDITIONAL_CERTS_FILE] "
              << "[--write-proto=PROTO_FILE] "
              << "[--write-cpp-root-store=CPP_FILE] "
              << "[--write-cpp-ev-roots=CPP_FILE] " << std::endl;
    return 1;
  }

  std::optional<RootStore> root_store =
      ReadTextRootStore(root_store_path, certs_path, additional_certs_path);
  if (!root_store) {
    return 1;
  }

  // TODO(crbug.com/40770548): Figure out how to use the serialized
  // proto to support component update.
  // components/resources/ssl/ssl_error_assistant/push_proto.py
  // does it through a GCS bucket (I think) so that might be an option.
  if (!proto_path.empty()) {
    std::string serialized;
    if (!root_store->SerializeToString(&serialized)) {
      LOG(ERROR) << "Error serializing root store proto"
                 << root_store->DebugString();
      return 1;
    }
    if (!base::WriteFile(proto_path, serialized)) {
      PLOG(ERROR) << "Error writing serialized proto root store";
      return 1;
    }
  }

  if (!root_store_cpp_path.empty() &&
      !WriteRootCppFile(*root_store, root_store_cpp_path)) {
    PLOG(ERROR) << "Error writing root store C++ include file";
    return 1;
  }
  if (!ev_roots_cpp_path.empty() &&
      !WriteEvCppFile(*root_store, ev_roots_cpp_path)) {
    PLOG(ERROR) << "Error writing EV roots C++ include file";
    return 1;
  }

  return 0;
}