/*
    * Copyright 2024 The Netty Project
    *
    * The Netty Project licenses this file to you 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:
    *
    *   https://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 io.netty.pkitesting;
  
  import io.netty.util.internal.EmptyArrays;
  
  import java.io.File;
  import java.io.IOException;
  import java.io.OutputStream;
  import java.nio.charset.StandardCharsets;
  import java.nio.file.Files;
  import java.nio.file.Path;
  import java.nio.file.StandardOpenOption;
  import java.security.KeyPair;
  import java.security.KeyStore;
  import java.security.KeyStoreException;
  import java.security.NoSuchAlgorithmException;
  import java.security.PrivateKey;
  import java.security.UnrecoverableKeyException;
  import java.security.cert.CertificateEncodingException;
  import java.security.cert.CertificateException;
  import java.security.cert.TrustAnchor;
  import java.security.cert.X509Certificate;
  import java.util.Arrays;
  import java.util.Base64;
  import java.util.List;
  import javax.net.ssl.KeyManagerFactory;
  import javax.net.ssl.TrustManager;
  import javax.net.ssl.TrustManagerFactory;
  
  import static java.util.Objects.requireNonNull;
  
  /**
   * A certificate bundle is a private key and a full certificate path, all the way to the root certificate.
   * The bundle offers ways of accessing these, and converting them into various representations.
   */
  public final class X509Bundle {
      private final X509Certificate[] certPath;
      private final X509Certificate root;
      private final KeyPair keyPair;
  
      /**
       * Construct a bundle from a given certificate path, root certificate, and {@link KeyPair}.
       * @param certPath The certificate path, starting with the leaf certificate.The path can end either with the
       * root certificate, or the intermediate certificate signed by the root certificate.
       * @param root The self-signed root certificate.
       * @param keyPair The key pair.
       */
      private X509Bundle(X509Certificate[] certPath, X509Certificate root, KeyPair keyPair) {
          requireNonNull(root, "root");
          requireNonNull(keyPair, "keyPair");
          if (certPath.length > 1 && certPath[certPath.length - 1].equals(root)) {
              this.certPath = Arrays.copyOf(certPath, certPath.length - 1);
          } else {
              this.certPath = certPath.clone();
          }
          this.root = root;
          this.keyPair = keyPair;
      }
  
      /**
       * Construct a bundle for a certificate authority.
       * @param root The self-signed root certificate.
       * @param keyPair The key pair.
       * @return The new bundle.
       */
      public static X509Bundle fromRootCertificateAuthority(X509Certificate root, KeyPair keyPair) {
          requireNonNull(root, "root");
          requireNonNull(keyPair, "keyPair");
          X509Bundle bundle = new X509Bundle(new X509Certificate[]{root}, root, keyPair);
          if (!bundle.isCertificateAuthority() || !bundle.isSelfSigned()) {
              throw new IllegalArgumentException("Given certificate is not a root CA certificate: " +
                      root.getSubjectX500Principal() + ", issued by " + root.getIssuerX500Principal());
          }
          return bundle;
      }
  
      /**
       * Construct a bundle from a given certificate path, root certificate, and {@link KeyPair}.
       * @param certPath The certificate path, starting with the leaf certificate.The path can end either with the
       * root certificate, or the intermediate certificate signed by the root certificate.
       * @param root The self-signed root certificate.
       * @param keyPair The key pair.
       */
      public static X509Bundle fromCertificatePath(
              X509Certificate[] certPath, X509Certificate root, KeyPair keyPair) {
         return new X509Bundle(certPath, root, keyPair);
     }
 
     /**
      * Get the leaf certificate of the bundle.
      * If this bundle is for a certificate authority, then this return the same as {@link #getRootCertificate()}.
      * @return The leaf certificate.
      */
     public X509Certificate getCertificate() {
         return certPath[0];
     }
 
     /**
      * Get the PEM encoded string of the {@linkplain #getCertificate() leaf certificate}.
      * @return The certificate PEM string.
      */
     public String getCertificatePEM() {
         return toCertPem(certPath[0]);
     }
 
     /**
      * Get the certificate path, starting with the leaf certificate up to but excluding the root certificate.
      * @return The certificate path.
      */
     public X509Certificate[] getCertificatePath() {
         return certPath.clone();
     }
 
     /**
      * Get the certificate path, starting with the leaf certificate up to and including the root certificate.
      * @return The certificate path, including the root certificate.
      */
     public X509Certificate[] getCertificatePathWithRoot() {
         X509Certificate[] path = Arrays.copyOf(certPath, certPath.length + 1);
         path[path.length - 1] = root;
         return path;
     }
 
     /**
      * Get the certificate path as a list, starting with the leaf certificate up to but excluding the root certificate.
      * @return The certificate path list.
      */
     public List<X509Certificate> getCertificatePathList() {
         return Arrays.asList(certPath);
     }
 
     /**
      * Get the {@linkplain #getCertificatePath() certificate path} as a PEM encoded string.
      * @return The PEM encoded certificate path.
      */
     public String getCertificatePathPEM() {
         return toCertPem(certPath);
     }
 
     /**
      * Get the key pair.
      * @return The key pair.
      */
     public KeyPair getKeyPair() {
         return keyPair;
     }
 
     /**
      * Get the root certificate that anchors the certificate path.
      * @return The root certificate.
      */
     public X509Certificate getRootCertificate() {
         return root;
     }
 
     /**
      * Get the {@linkplain #getRootCertificate() root certificate} as a PEM encoded string.
      * @return The PEM encoded root certificate.
      */
     public String getRootCertificatePEM() {
         return toCertPem(root);
     }
 
     private static String toCertPem(X509Certificate... certs) {
         Base64.Encoder encoder = getMimeEncoder();
         StringBuilder sb = new StringBuilder();
         for (X509Certificate cert : certs) {
             sb.append("-----BEGIN CERTIFICATE-----\r\n");
             try {
                 sb.append(encoder.encodeToString(cert.getEncoded()));
             } catch (CertificateEncodingException e) {
                 throw new IllegalStateException(e);
             }
             sb.append("\r\n-----END CERTIFICATE-----\r\n");
         }
         return sb.toString();
     }
 
     /**
      * Get the private key as a PEM encoded PKCS#8 string.
      * @return The private key in PKCS#8 and PEM encoded string.
      */
     public String getPrivateKeyPEM() {
         Base64.Encoder encoder = getMimeEncoder();
         StringBuilder sb = new StringBuilder();
         sb.append("-----BEGIN PRIVATE KEY-----\r\n");
         PrivateKey privateKey = keyPair.getPrivate();
         sb.append(encoder.encodeToString(privateKey.getEncoded()));
         sb.append("\r\n-----END PRIVATE KEY-----\r\n");
         return sb.toString();
     }
 
     private static Base64.Encoder getMimeEncoder() {
         return Base64.getMimeEncoder(64, new byte[]{ '\r', '\n' });
     }
 
     /**
      * Get the root certificate as a new {@link TrustAnchor} object.
      * Note that {@link TrustAnchor} instance have object identity, so if this method is called twice,
      * the two trust anchors will not be equal to each other.
      * @return A new {@link TrustAnchor} instance containing the root certificate.
      */
     public TrustAnchor getTrustAnchor() {
         return new TrustAnchor(root, root.getExtensionValue(CertificateBuilder.OID_X509_NAME_CONSTRAINTS));
     }
 
     /**
      * Query if this bundle is for a certificate authority root certificate.
      * @return {@code true} if the {@linkplain #getCertificate() leaf certificate} is a certificate authority,
      * otherwise {@code false}.
      */
     public boolean isCertificateAuthority() {
         return certPath[0].getBasicConstraints() != -1;
     }
 
     /**
      * Query if this bundle is for a self-signed certificate.
      * @return {@code true} if the {@linkplain #getCertificate() leaf certificate} is self-signed.
      */
     public boolean isSelfSigned() {
         X509Certificate leaf = certPath[0];
         return certPath.length == 1 &&
                 leaf.getSubjectX500Principal().equals(leaf.getIssuerX500Principal()) &&
                 Arrays.equals(leaf.getSubjectUniqueID(), leaf.getIssuerUniqueID());
     }
 
     /**
      * Create a {@link TrustManager} instance that trusts the root certificate in this bundle.
      * @return The new {@link TrustManager}.
      */
     public TrustManager toTrustManager() {
         TrustManagerFactory tmf = toTrustManagerFactory();
         return tmf.getTrustManagers()[0];
     }
 
     /**
      * Create {@link TrustManagerFactory} instance that trusts the root certificate in this bundle.
      * <p>
      * The trust manager factory will use the {@linkplain TrustManagerFactory#getDefaultAlgorithm() default algorithm}.
      *
      * @return The new {@link TrustManagerFactory}.
      */
     public TrustManagerFactory toTrustManagerFactory() {
         return toTrustManagerFactory(TrustManagerFactory.getDefaultAlgorithm());
     }
 
     /**
      * Create {@link TrustManagerFactory} instance that trusts the root certificate in this bundle,
      * with the given algorithm.
      *
      * @return The new {@link TrustManagerFactory}.
      */
     public TrustManagerFactory toTrustManagerFactory(String algorithm) {
         try {
             TrustManagerFactory tmf = TrustManagerFactory.getInstance(algorithm);
             tmf.init(toKeyStore(EmptyArrays.EMPTY_CHARS));
             return tmf;
         } catch (NoSuchAlgorithmException e) {
             throw new AssertionError("Default TrustManagerFactory algorithm was not available.", e);
         } catch (KeyStoreException e) {
             throw new IllegalStateException("Failed to initialize TrustManagerFactory with KeyStore.", e);
         }
     }
 
     /**
      * Create a {@link KeyStore} with the contents of this bundle.
      * The root certificate will be a trusted root in the key store.
      * If this bundle has a {@linkplain #getPrivateKeyPEM() private key},
      * then the private key and certificate path will also be added to the key store.
      * <p>
      * The key store will use the PKCS#12 format.
      *
      * @param keyEntryPassword The password used to encrypt the private key entry in the key store.
      * @return The key store.
      * @throws KeyStoreException If an error occurred when adding entries to the key store.
      */
     public KeyStore toKeyStore(char[] keyEntryPassword) throws KeyStoreException {
         return toKeyStore("PKCS12", keyEntryPassword);
     }
 
     /**
      * Create a {@link KeyStore} with the contents of this bundle.
      * The root certificate will be a trusted root in the key store.
      * If this bundle has a {@linkplain #getPrivateKeyPEM() private key},
      * then the private key and certificate path will also be added to the key store.
      * <p>
      * The key store will use the format defined by the given algorithm.
      *
      * @param keyEntryPassword The password used to encrypt the private key entry in the key store.
      * @return The key store.
      * @throws KeyStoreException If an error occurred when adding entries to the key store.
      */
     public KeyStore toKeyStore(String algorithm, char[] keyEntryPassword) throws KeyStoreException {
         KeyStore keyStore;
         try {
             keyStore = KeyStore.getInstance(algorithm);
             keyStore.load(null, null);
         } catch (IOException | NoSuchAlgorithmException | CertificateException e) {
             throw new KeyStoreException("Failed to initialize '" + algorithm + "' KeyStore.", e);
         }
         keyStore.setCertificateEntry("1", root);
         if (keyPair.getPrivate() != null) {
             keyStore.setKeyEntry("2", MLDSASeedPrivateKey.unwrap(keyPair.getPrivate()), keyEntryPassword, certPath);
         }
         return keyStore;
     }
 
     /**
      * Create a temporary PKCS#12 file with the {@linkplain #toKeyStore(char[]) key store} of this bundle.
      * The temporary file is automatically deleted when the JVM terminates normally.
      * @param password The password used both to encrypt the private key in the key store,
      * and to protect the key store itself.
      * @return The {@link File} object with the path to the PKCS#12 key store.
      * @throws Exception If something went wrong with creating the key store file.
      */
     public File toTempKeyStoreFile(char[] password) throws Exception {
         return toTempKeyStoreFile(password, password);
     }
 
     /**
      * Create a temporary PKCS#12 file with the {@linkplain #toKeyStore(char[]) key store} of this bundle.
      * The temporary file is automatically deleted when the JVM terminates normally.
      * @param pkcs12Password The password used to encrypt the PKCS#12 file.
      * @param keyEntryPassword The password used to encrypt the private key entry in the PKCS#12 file.
      * @return The {@link File} object with the path to the PKCS#12 key store.
      * @throws Exception If something went wrong with creating the key store file.
      */
     public File toTempKeyStoreFile(char[] pkcs12Password, char[] keyEntryPassword) throws Exception {
         KeyStore keyStore = toKeyStore(keyEntryPassword);
         Path tempFile = Files.createTempFile("ks", ".p12");
         try (OutputStream out = Files.newOutputStream(tempFile, StandardOpenOption.WRITE)) {
             keyStore.store(out, pkcs12Password);
         }
         File file = tempFile.toFile();
         file.deleteOnExit();
         return file;
     }
 
     /**
      * Create a temporary PEM file with the {@linkplain #getRootCertificate() root certificate} of this bundle.
      * The temporary file is automatically deleted when the JVM terminates normally.
      * @return The {@link File} object with the path to the trust root PEM file.
      * @throws IOException If an IO error occurred when creating the trust root file.
      */
     public File toTempRootCertPem() throws IOException {
         return createTempPemFile(getRootCertificatePEM(), "ca");
     }
 
     /**
      * Create a temporary PEM file with the {@linkplain #getCertificatePath() certificate chain} of this bundle.
      * The temporary file is automatically deleted when the JVM terminates normally.
      * @return The {@link File} object with the path to the certificate chain PEM file.
      * @throws IOException If an IO error occurred when creating the certificate chain file.
      */
     public File toTempCertChainPem() throws IOException {
         return createTempPemFile(getCertificatePathPEM(), "chain");
     }
 
     /**
      * Create a temporary PEM file with the {@linkplain #getPrivateKeyPEM() private key} of this bundle.
      * The temporary file is automatically deleted when the JVM terminates normally.
      * @return The {@link File} object with the path to the private key PEM file.
      * @throws IOException If an IO error occurred when creating the private key file.
      */
     public File toTempPrivateKeyPem() throws IOException {
         return createTempPemFile(getPrivateKeyPEM(), "key");
     }
 
     private static File createTempPemFile(String pem, String filePrefix) throws IOException {
         Path tempFile = Files.createTempFile(filePrefix, ".pem");
         try (OutputStream out = Files.newOutputStream(tempFile, StandardOpenOption.WRITE)) {
             out.write(pem.getBytes(StandardCharsets.ISO_8859_1));
         }
         File file = tempFile.toFile();
         file.deleteOnExit();
         return file;
     }
 
     /**
      * Create a {@link KeyManagerFactory} from this bundle.
      * <p>
      * The {@link KeyManagerFactory} will use the
      * {@linkplain KeyManagerFactory#getDefaultAlgorithm() default algorithm}.
      *
      * @return The new {@link KeyManagerFactory}.
      * @throws KeyStoreException If there was a problem creating or initializing the key store.
      * @throws UnrecoverableKeyException If the private key could not be recovered,
      * for instance if this bundle is a {@linkplain #isCertificateAuthority() certificate authority}.
      * @throws NoSuchAlgorithmException If the key manager factory algorithm is not supported by the current
      * security provider.
      */
     public KeyManagerFactory toKeyManagerFactory()
             throws KeyStoreException, UnrecoverableKeyException, NoSuchAlgorithmException {
         return toKeyManagerFactory(KeyManagerFactory.getDefaultAlgorithm());
     }
 
     /**
      * Create a {@link KeyManagerFactory} from this bundle, using the given algorithm.
      *
      * @return The new {@link KeyManagerFactory}.
      * @throws KeyStoreException If there was a problem creating or initializing the key store.
      * @throws UnrecoverableKeyException If the private key could not be recovered,
      * for instance if this bundle is a {@linkplain #isCertificateAuthority() certificate authority}.
      * @throws NoSuchAlgorithmException If the key manager factory algorithm is not supported by the current
      * security provider.
      */
     public KeyManagerFactory toKeyManagerFactory(String algorithm)
             throws KeyStoreException, UnrecoverableKeyException, NoSuchAlgorithmException {
         KeyManagerFactory kmf;
         try {
             kmf = KeyManagerFactory.getInstance(algorithm);
         } catch (NoSuchAlgorithmException e) {
             throw new AssertionError("Default KeyManagerFactory algorithm was not available.", e);
         }
         kmf.init(toKeyStore(EmptyArrays.EMPTY_CHARS), EmptyArrays.EMPTY_CHARS);
         return kmf;
     }
 }