// Copyright 2020 The Nomulus Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package google.registry.flows.certs;

import static com.google.common.base.Preconditions.checkArgument;
import static google.registry.util.DateTimeUtils.START_INSTANT;
import static google.registry.util.DateTimeUtils.plusDays;
import static google.registry.util.DateTimeUtils.toLocalDate;
import static java.nio.charset.StandardCharsets.UTF_8;

import com.google.common.collect.ImmutableSet;
import com.google.common.collect.ImmutableSortedMap;
import google.registry.config.RegistryConfig.Config;
import google.registry.util.Clock;
import jakarta.inject.Inject;
import java.io.ByteArrayInputStream;
import java.io.StringWriter;
import java.security.PublicKey;
import java.security.cert.CertificateException;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.security.interfaces.ECPublicKey;
import java.security.interfaces.RSAPublicKey;
import java.time.Instant;
import java.time.temporal.ChronoUnit;
import java.util.stream.Collectors;
import org.bouncycastle.jcajce.provider.asymmetric.util.EC5Util;
import org.bouncycastle.jce.ECNamedCurveTable;
import org.bouncycastle.jce.spec.ECNamedCurveParameterSpec;
import org.bouncycastle.jce.spec.ECParameterSpec;
import org.bouncycastle.openssl.jcajce.JcaMiscPEMGenerator;
import org.bouncycastle.util.io.pem.PemObjectGenerator;
import org.bouncycastle.util.io.pem.PemWriter;

/** A utility to check that a given certificate meets our requirements */
public class CertificateChecker {

  private final ImmutableSortedMap<Instant, Integer> maxValidityLengthSchedule;
  private final long expirationWarningDays;
  private final int minimumRsaKeyLength;
  private final Clock clock;
  private final ImmutableSet<String> allowedEcdsaCurves;
  private final long expirationWarningIntervalDays;

  /**
   * Constructs a CertificateChecker instance with the specified configuration parameters.
   *
   * <p>The max validity length schedule is a sorted map of {@link Instant} to {@link Integer}
   * entries representing a maximum validity period for certificates issued on or after that date.
   * The first entry must have a key of {@link google.registry.util.DateTimeUtils#START_INSTANT},
   * such that every possible date has an applicable max validity period. Since security
   * requirements tighten over time, the max validity periods will be decreasing as the date
   * increases.
   *
   * <p>The validity length schedule used by all major Web browsers as of 2020Q4 would be
   * represented as:
   *
   * <pre>
   *   ImmutableSortedMap.of(
   *     START_INSTANT, 825,
   *     Instant.parse("2020-09-01T00:00:00Z"), 398
   *   );
   * </pre>
   */
  @Inject
  public CertificateChecker(
      @Config("maxValidityDaysSchedule")
          ImmutableSortedMap<Instant, Integer> maxValidityDaysSchedule,
      @Config("expirationWarningDays") long expirationWarningDays,
      @Config("expirationWarningIntervalDays") long expirationWarningIntervalDays,
      @Config("minimumRsaKeyLength") int minimumRsaKeyLength,
      @Config("allowedEcdsaCurves") ImmutableSet<String> allowedEcdsaCurves,
      Clock clock) {
    checkArgument(
        maxValidityDaysSchedule.containsKey(START_INSTANT),
        "Max validity length schedule must contain an entry for START_INSTANT");
    this.maxValidityLengthSchedule = maxValidityDaysSchedule;
    this.expirationWarningDays = expirationWarningDays;
    this.minimumRsaKeyLength = minimumRsaKeyLength;
    this.allowedEcdsaCurves = allowedEcdsaCurves;
    this.expirationWarningIntervalDays = expirationWarningIntervalDays;
    this.clock = clock;
  }

  private static long getValidityLengthInDays(X509Certificate certificate) {
    return ChronoUnit.DAYS.between(
        toLocalDate(certificate.getNotBefore().toInstant()),
        toLocalDate(certificate.getNotAfter().toInstant()));
  }

  /** Checks if the curve used for a public key is in the list of acceptable curves. */
  private static boolean checkCurveName(PublicKey key, ImmutableSet<String> allowedEcdsaCurves) {
    ECParameterSpec params;
    // These 2 different instances of PublicKey need to be handled separately since their OIDs are
    // encoded differently. More details on this can be found at
    // https://stackoverflow.com/questions/49895713/how-to-find-the-matching-curve-name-from-an-ecpublickey.
    if (key instanceof ECPublicKey ecKey) {
      params = EC5Util.convertSpec(ecKey.getParams());
    } else if (key instanceof org.bouncycastle.jce.interfaces.ECPublicKey ecKey) {
      params = ecKey.getParameters();
    } else {
      throw new IllegalArgumentException("Unrecognized instance of PublicKey.");
    }
    return allowedEcdsaCurves.stream()
        .anyMatch(
            curve -> {
              ECNamedCurveParameterSpec cParams = ECNamedCurveTable.getParameterSpec(curve);
              return cParams.getN().equals(params.getN())
                  && cParams.getH().equals(params.getH())
                  && cParams.getCurve().equals(params.getCurve())
                  && cParams.getG().equals(params.getG());
            });
  }

  /**
   * Checks the given certificate string for violations and throws an exception if any violations
   * exist.
   */
  public void validateCertificate(String certificateString) throws InsecureCertificateException {
    handleCertViolations(checkCertificate(certificateString));
  }

  /**
   * Checks the given certificate string for violations and throws an exception if any violations
   * exist.
   */
  public void validateCertificate(X509Certificate certificate) throws InsecureCertificateException {
    handleCertViolations(checkCertificate(certificate));
  }

  private void handleCertViolations(ImmutableSet<CertificateViolation> violations)
      throws InsecureCertificateException {
    if (!violations.isEmpty()) {
      String displayMessages =
          violations.stream()
              .map(this::getViolationDisplayMessage)
              .collect(Collectors.joining("\n"));
      throw new InsecureCertificateException(violations, displayMessages);
    }
  }

  /**
   * Checks a given certificate for violations and returns a list of all the violations the
   * certificate has.
   */
  public ImmutableSet<CertificateViolation> checkCertificate(X509Certificate certificate) {
    ImmutableSet.Builder<CertificateViolation> violations = new ImmutableSet.Builder<>();

    // Check if currently in validity period
    Instant now = clock.now();
    if (certificate.getNotAfter().toInstant().isBefore(now)) {
      violations.add(CertificateViolation.EXPIRED);
    } else if (certificate.getNotBefore().toInstant().isAfter(now)) {
      violations.add(CertificateViolation.NOT_YET_VALID);
    }

    // Check validity period length
    long maxValidityDays =
        maxValidityLengthSchedule.floorEntry(certificate.getNotBefore().toInstant()).getValue();
    if (getValidityLengthInDays(certificate) > maxValidityDays) {
      violations.add(CertificateViolation.VALIDITY_LENGTH_TOO_LONG);
    }

    // Check key strengths
    PublicKey key = certificate.getPublicKey();
    if (key.getAlgorithm().equals("RSA")) {
      RSAPublicKey rsaPublicKey = (RSAPublicKey) key;
      if (rsaPublicKey.getModulus().bitLength() < minimumRsaKeyLength) {
        violations.add(CertificateViolation.RSA_KEY_LENGTH_TOO_SHORT);
      }
    } else if (key.getAlgorithm().equals("EC")) {
      if (!checkCurveName(key, allowedEcdsaCurves)) {
        violations.add(CertificateViolation.INVALID_ECDSA_CURVE);
      }
    } else {
      violations.add(CertificateViolation.ALGORITHM_CONSTRAINED);
    }
    return violations.build();
  }

  /**
   * Converts a given string to a certificate and checks it for violations, returning a list of all
   * the violations the certificate has.
   */
  public ImmutableSet<CertificateViolation> checkCertificate(String certificateString) {
    return checkCertificate(getCertificate(certificateString));
  }

  /** Converts the given string to a certificate object. */
  public X509Certificate getCertificate(String certificateStr) {
    X509Certificate certificate;
    try {
      certificate =
          (X509Certificate)
              CertificateFactory.getInstance("X509")
                  .generateCertificate(new ByteArrayInputStream(certificateStr.getBytes(UTF_8)));
    } catch (CertificateException e) {
      throw new IllegalArgumentException(
          String.format("Unable to read given certificate %s", certificateStr), e);
    }
    return certificate;
  }

  /** Serializes the certificate object to a certificate string. */
  public String serializeCertificate(X509Certificate certificate) throws Exception {
    StringWriter sw = new StringWriter();
    try (PemWriter pw = new PemWriter(sw)) {
      PemObjectGenerator generator = new JcaMiscPEMGenerator(certificate);
      pw.writeObject(generator);
    }
    return sw.toString();
  }

  /** Returns whether the client should receive a notification email. */
  public boolean shouldReceiveExpiringNotification(
      Instant lastExpiringNotificationSentDate, String certificateStr) {
    X509Certificate certificate = getCertificate(certificateStr);
    Instant now = clock.now();
    // the expiration date is one day after lastValidDate
    Instant lastValidDate = certificate.getNotAfter().toInstant();
    if (lastValidDate.isBefore(now)) {
      return false;
    }
    /*
     * Client should receive a notification if:
     *    1) client has never received notification (lastExpiringNotificationSentDate is initially
     *    set to START_INSTANT) and the certificate has entered the expiring period, OR
     *    2) client has received notification but the interval between now and
     *    lastExpiringNotificationSentDate is greater than expirationWarningIntervalDays.
     */
    return !lastValidDate.isAfter(plusDays(now, expirationWarningDays))
        && (lastExpiringNotificationSentDate.equals(START_INSTANT)
            || !plusDays(lastExpiringNotificationSentDate, expirationWarningIntervalDays)
                .isAfter(now));
  }

  private String getViolationDisplayMessage(CertificateViolation certificateViolation) {
    // Yes, we'd rather do this as an instance method on the CertificateViolation enum itself, but
    // we can't because we need access to configuration (injected as instance variables) which you
    // can't get in a static enum context.
    return switch (certificateViolation) {
      case EXPIRED -> "Certificate is expired.";
      case NOT_YET_VALID -> "Certificate start date is in the future.";
      case ALGORITHM_CONSTRAINED -> "Certificate key algorithm must be RSA or ECDSA.";
      case RSA_KEY_LENGTH_TOO_SHORT ->
          String.format(
              "RSA key length is too short; the minimum allowed length is %d bits.",
              this.minimumRsaKeyLength);
      case VALIDITY_LENGTH_TOO_LONG ->
          String.format(
              "Certificate validity period is too long; it must be less than or equal to %d days.",
              this.maxValidityLengthSchedule.lastEntry().getValue());
      case INVALID_ECDSA_CURVE ->
          String.format("The ECDSA key must use one of these algorithms: %s", allowedEcdsaCurves);
    };
  }

  /**
   * The type of violation a certificate has based on the certificate requirements
   * (go/registry-proxy-security).
   */
  public enum CertificateViolation {
    EXPIRED,
    NOT_YET_VALID,
    VALIDITY_LENGTH_TOO_LONG,
    RSA_KEY_LENGTH_TOO_SHORT,
    ALGORITHM_CONSTRAINED,
    INVALID_ECDSA_CURVE;

    /**
     * Gets a suitable end-user-facing display message for this particular certificate violation.
     *
     * <p>Note that the {@link CertificateChecker} instance must be passed in because it contains
     * configuration values (e.g., minimum RSA key length) that go into the error message text.
     */
    public String getDisplayMessage(CertificateChecker certificateChecker) {
      return certificateChecker.getViolationDisplayMessage(this);
    }
  }

  /** Exception to throw when a certificate has security violations. */
  public static class InsecureCertificateException extends Exception {
    ImmutableSet<CertificateViolation> violations;

    InsecureCertificateException(ImmutableSet<CertificateViolation> violations, String message) {
      super(message);
      this.violations = violations;
    }
  }
}
