// Copyright 2017 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.networking.module;

import static com.google.common.base.Preconditions.checkState;
import static com.google.common.base.Suppliers.memoizeWithExpiration;
import static com.google.common.collect.ImmutableList.toImmutableList;
import static java.nio.charset.StandardCharsets.UTF_8;
import static java.util.concurrent.TimeUnit.SECONDS;

import com.google.common.base.Suppliers;
import com.google.common.collect.ImmutableList;
import dagger.Lazy;
import dagger.Module;
import dagger.Provides;
import google.registry.util.Clock;
import google.registry.util.SelfSignedCaCertificate;
import jakarta.inject.Named;
import jakarta.inject.Provider;
import jakarta.inject.Qualifier;
import jakarta.inject.Singleton;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.security.PrivateKey;
import java.security.Security;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;
import java.time.Duration;
import java.util.function.Function;
import java.util.function.Supplier;
import org.bouncycastle.asn1.pkcs.PrivateKeyInfo;
import org.bouncycastle.cert.X509CertificateHolder;
import org.bouncycastle.cert.jcajce.JcaX509CertificateConverter;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.openssl.PEMException;
import org.bouncycastle.openssl.PEMKeyPair;
import org.bouncycastle.openssl.PEMParser;
import org.bouncycastle.openssl.jcajce.JcaPEMKeyConverter;

/**
 * Dagger module that provides bindings needed to inject a certificate chain the corresponding
 * private key.
 *
 * <p>There are several ways ({@link Mode}s) that the certificate/key can be provided (explained
 * later), and the user of this module needs to provide the binding to the {@link Mode} in some
 * other module to the component.
 *
 * <p>The production certificates and private key are stored in a .pem file that is encrypted by
 * Cloud KMS. The .pem file can be generated by concatenating the .crt certificate files on the
 * chain and the .key private file.
 *
 * <p>The production certificates in the .pem file must be stored in order, where the next
 * certificate's subject is the previous certificate's issuer.
 *
 * <p>When running the proxy locally or in test, a self signed certificate is used.
 *
 * @see <a href="https://cloud.google.com/kms/">Cloud Key Management Service</a>
 */
@Module
public final class CertificateSupplierModule {

  static {
    Security.addProvider(new BouncyCastleProvider());
  }

  private CertificateSupplierModule() {}

  public enum Mode {

    /**
     * The certificate chain and the private key are stored in a single PEM file.
     *
     * <p>Both certificates and private key are Base64 encoded and the single PEM file is generated
     * by concatenating each ASCII-armored block. The private key can appear anywhere in the list
     * but the certificates must appear in order, i. e. the next certificate's subject is the
     * previous certificates's issuer.
     *
     * @see <a href="https://en.wikipedia.org/wiki/Privacy-Enhanced_Mail">Privacy-Enhanced Email</a>
     */
    PEM_FILE,

    P12_FILE,

    /**
     * A single certificate/private key pair is generated in place and self signed. Used in tests.
     */
    SELF_SIGNED
  }

  @Qualifier
  @interface PemFile {}

  @Qualifier
  private @interface P12File {}

  @Qualifier
  private @interface SelfSigned {}

  /**
   * Select specific type from a given {@link ImmutableList} and convert them using the converter.
   *
   * @param objects the {@link ImmutableList} to filter from.
   * @param clazz the class to filter.
   * @param converter the converter function to act on the items in the filtered list.
   */
  private static <T, E> ImmutableList<E> filterAndConvert(
      ImmutableList<Object> objects, Class<T> clazz, Function<T, E> converter) {
    return objects.stream()
        .filter(clazz::isInstance)
        .map(clazz::cast)
        .map(converter)
        .collect(toImmutableList());
  }

  @Singleton
  @Provides
  static Supplier<PrivateKey> providePrivateKeySupplier(
      Mode mode,
      @PemFile Lazy<Supplier<PrivateKey>> pemPrivateKeySupplier,
      @P12File Lazy<Supplier<PrivateKey>> p12PrivateKeySupplier,
      @SelfSigned Lazy<Supplier<PrivateKey>> selfSignedPrivateKeySupplier) {
    return switch (mode) {
      case PEM_FILE -> pemPrivateKeySupplier.get();
      case P12_FILE -> p12PrivateKeySupplier.get();
      case SELF_SIGNED -> selfSignedPrivateKeySupplier.get();
      default -> throw new RuntimeException("Certificate provider mode exhausted.");
    };
  }

  @Singleton
  @Provides
  static Supplier<ImmutableList<X509Certificate>> provideCertificatesSupplier(
      Mode mode,
      @PemFile Lazy<Supplier<ImmutableList<X509Certificate>>> pemCertificatesSupplier,
      @P12File Lazy<Supplier<ImmutableList<X509Certificate>>> p12CertificatesSupplier,
      @SelfSigned Lazy<Supplier<ImmutableList<X509Certificate>>> selfSignedCertificatesSupplier) {
    return switch (mode) {
      case PEM_FILE -> pemCertificatesSupplier.get();
      case P12_FILE -> p12CertificatesSupplier.get();
      case SELF_SIGNED -> selfSignedCertificatesSupplier.get();
      default -> throw new RuntimeException("Certificate provider mode exhausted.");
    };
  }

  @Singleton
  @Provides
  static SelfSignedCaCertificate provideSelfSignedCertificate(Clock clock) {
    try {
      return SelfSignedCaCertificate.create(clock);
    } catch (Exception e) {
      throw new RuntimeException(e);
    }
  }

  @Singleton
  @Provides
  @SelfSigned
  static Supplier<PrivateKey> provideSelfSignedPrivateKeySupplier(SelfSignedCaCertificate ssc) {
    return Suppliers.ofInstance(ssc.key());
  }

  @Singleton
  @Provides
  @SelfSigned
  static Supplier<ImmutableList<X509Certificate>> provideSelfSignedCertificatesSupplier(
      SelfSignedCaCertificate ssc) {
    return Suppliers.ofInstance(ImmutableList.of(ssc.cert()));
  }

  @Provides
  @PemFile
  static ImmutableList<Object> providePemObjects(@Named("pemBytes") byte[] pemBytes) {
    PEMParser pemParser =
        new PEMParser(new InputStreamReader(new ByteArrayInputStream(pemBytes), UTF_8));
    ImmutableList.Builder<Object> listBuilder = new ImmutableList.Builder<>();
    Object obj;
    // PEMParser returns an object (private key, certificate, etc) each time readObject() is called,
    // until no more object is to be read from the file.
    while (true) {
      try {
        obj = pemParser.readObject();
        if (obj == null) {
          break;
        } else {
          listBuilder.add(obj);
        }
      } catch (IOException e) {
        throw new RuntimeException("Cannot parse PEM file correctly.", e);
      }
    }
    return listBuilder.build();
  }

  // This binding should not be used directly. Use the supplier binding instead.
  @Provides
  @PemFile
  static PrivateKey providePemPrivateKey(@PemFile ImmutableList<Object> pemObjects) {
    JcaPEMKeyConverter converter = new JcaPEMKeyConverter().setProvider("BC");
    Function<PEMKeyPair, PrivateKey> pkcs1PrivateKeyConverter =
        pemKeyPair -> {
          try {
            return converter.getKeyPair(pemKeyPair).getPrivate();
          } catch (PEMException e) {
            throw new RuntimeException(
                String.format("Error converting PKCS#1 private key: %s", pemKeyPair), e);
          }
        };
    Function<PrivateKeyInfo, PrivateKey> pkcs8PrivateKeyConverter =
        privateKeyInfo -> {
          try {
            return converter.getPrivateKey(privateKeyInfo);
          } catch (PEMException e) {
            throw new RuntimeException(
                String.format("Error converting PKCS#8 private key: %s", privateKeyInfo), e);
          }
        };
    ImmutableList<PrivateKey> privateKeys =
        ImmutableList.<PrivateKey>builder()
            .addAll(filterAndConvert(pemObjects, PEMKeyPair.class, pkcs1PrivateKeyConverter))
            .addAll(filterAndConvert(pemObjects, PrivateKeyInfo.class, pkcs8PrivateKeyConverter))
            .build();
    checkState(
        privateKeys.size() == 1,
        "The pem file must contain exactly one private key, but %s keys are found",
        privateKeys.size());
    return privateKeys.get(0);
  }

  // This binding should not be used directly. Use the supplier binding instead.
  @Provides
  @PemFile
  static ImmutableList<X509Certificate> providePemCertificates(
      @PemFile ImmutableList<Object> pemObject) {
    JcaX509CertificateConverter converter = new JcaX509CertificateConverter().setProvider("BC");
    Function<X509CertificateHolder, X509Certificate> certificateConverter =
        certificateHolder -> {
          try {
            return converter.getCertificate(certificateHolder);
          } catch (CertificateException e) {
            throw new RuntimeException(
                String.format("Error converting certificate: %s", certificateHolder), e);
          }
        };
    ImmutableList<X509Certificate> certificates =
        filterAndConvert(pemObject, X509CertificateHolder.class, certificateConverter);
    checkState(!certificates.isEmpty(), "No certificates found in the pem file");
    X509Certificate lastCert = null;
    for (X509Certificate cert : certificates) {
      if (lastCert != null) {
        checkState(
            lastCert.getIssuerX500Principal().equals(cert.getSubjectX500Principal()),
            "Certificate chain error:\n%s\nis not signed by\n%s",
            lastCert,
            cert);
      }
      lastCert = cert;
    }
    return certificates;
  }

  @Singleton
  @Provides
  @PemFile
  static Supplier<PrivateKey> providePemPrivateKeySupplier(
      @PemFile Provider<PrivateKey> privateKeyProvider,
      @Named("remoteCertCachingDuration") Duration cachingDuration) {
    return memoizeWithExpiration(privateKeyProvider::get, cachingDuration.toSeconds(), SECONDS);
  }

  @Singleton
  @Provides
  @PemFile
  static Supplier<ImmutableList<X509Certificate>> providePemCertificatesSupplier(
      @PemFile Provider<ImmutableList<X509Certificate>> certificatesProvider,
      @Named("remoteCertCachingDuration") Duration cachingDuration) {
    return memoizeWithExpiration(certificatesProvider::get, cachingDuration.toSeconds(), SECONDS);
  }

  // TODO(jianglai): Implement P12 supplier or convert the file to PEM format.
  @Singleton
  @Provides
  @P12File
  static Supplier<PrivateKey> provideP12PrivateKeySupplier() {
    return Suppliers.ofInstance(null);
  }

  // TODO(jianglai): Implement P12 supplier or convert the file to PEM format.
  @Singleton
  @Provides
  @P12File
  static Supplier<ImmutableList<X509Certificate>> provideP12CertificatesSupplier() {
    return Suppliers.ofInstance(null);
  }
}
