content: /*
   * Licensed to the Apache Software Foundation (ASF) under one or more
   * contributor license agreements.  See the NOTICE file distributed with
   * this work for additional information regarding copyright ownership.
   * The ASF 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
   *
   *    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 org.apache.celeborn.common.network.ssl;
  
  import static org.apache.celeborn.common.network.TestHelper.getResourceAsAbsolutePath;
  
  import java.io.File;
  import java.io.FileOutputStream;
  import java.io.IOException;
  import java.math.BigInteger;
  import java.nio.file.Files;
  import java.nio.file.StandardCopyOption;
  import java.security.*;
  import java.security.cert.Certificate;
  import java.security.cert.X509Certificate;
  import java.util.*;
  import java.util.stream.Stream;
  
  import org.bouncycastle.asn1.x500.X500Name;
  import org.bouncycastle.asn1.x509.BasicConstraints;
  import org.bouncycastle.asn1.x509.Extension;
  import org.bouncycastle.asn1.x509.GeneralName;
  import org.bouncycastle.asn1.x509.GeneralNames;
  import org.bouncycastle.cert.X509v3CertificateBuilder;
  import org.bouncycastle.cert.jcajce.JcaX509CertificateConverter;
  import org.bouncycastle.cert.jcajce.JcaX509CertificateHolder;
  import org.bouncycastle.cert.jcajce.JcaX509v3CertificateBuilder;
  import org.bouncycastle.jce.provider.BouncyCastleProvider;
  import org.bouncycastle.operator.ContentSigner;
  import org.bouncycastle.operator.jcajce.JcaContentSignerBuilder;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  public class SslSampleConfigs {
  
    private static final Logger LOG = LoggerFactory.getLogger(SslSampleConfigs.class);
  
    static {
      Security.addProvider(new BouncyCastleProvider());
    }
  
    public static final String DEFAULT_KEY_STORE_PATH = getResourceAsAbsolutePath("/ssl/server.jks");
    public static final String SECOND_KEY_STORE_PATH =
        getResourceAsAbsolutePath("/ssl/server_another.jks");
  
    // trust store has ca's for both keys.
    public static final String TRUST_STORE_PATH = getResourceAsAbsolutePath("/ssl/truststore.jks");
  
    // this is a trust store which does not have either the primary or second cert's ca
    public static final String TRUST_STORE_WITHOUT_CA =
        getResourceAsAbsolutePath("/ssl/truststore-without-ca.jks");
  
    public static Map<String, String> createDefaultConfigMapForModule(String module) {
      return createConfigMapForModule(module, true);
    }
  
    public static Map<String, String> createAnotherConfigMapForModule(String module) {
      return createConfigMapForModule(module, false);
    }
  
    public static Map<String, String> createAutoSslConfigForModule(String module) {
      Map<String, String> confMap = new HashMap<>();
      confMap.put("celeborn.ssl." + module + ".enabled", "true");
      confMap.put("celeborn.ssl." + module + ".protocol", "TLSv1.2");
      confMap.put("celeborn.ssl." + module + ".autoSslEnabled", "true");
      return confMap;
    }
  
    private static Map<String, String> createConfigMapForModule(String module, boolean forDefault) {
      Map<String, String> confMap = new HashMap<>();
      confMap.put("celeborn.ssl." + module + ".enabled", "true");
      confMap.put("celeborn.ssl." + module + ".trustStoreReloadingEnabled", "false");
      confMap.put("celeborn.ssl." + module + ".trustStoreReloadIntervalMs", "10000");
      if (forDefault) {
        confMap.put("celeborn.ssl." + module + ".keyStore", DEFAULT_KEY_STORE_PATH);
      } else {
        confMap.put("celeborn.ssl." + module + ".keyStore", SECOND_KEY_STORE_PATH);
      }
      confMap.put("celeborn.ssl." + module + ".keyStorePassword", "password");
      confMap.put("celeborn.ssl." + module + ".keyPassword", "password");
      confMap.put("celeborn.ssl." + module + ".privateKeyPassword", "password");
      confMap.put("celeborn.ssl." + module + ".protocol", "TLSv1.2");
      confMap.put("celeborn.ssl." + module + ".trustStore", TRUST_STORE_PATH);
      confMap.put("celeborn.ssl." + module + ".trustStorePassword", "password");
      return confMap;
    }
  
    public static void createTrustStore(
        File trustStore, String password, String alias, Certificate cert)
        throws GeneralSecurityException, IOException {
      KeyStore ks = createEmptyKeyStore();
      ks.setCertificateEntry(alias, cert);
      saveKeyStore(ks, trustStore, password);
    }
  
    /** Creates a keystore with multiple keys and saves it to a file. */
    public static <T extends Certificate> void createTrustStore(
        File trustStore, String password, Map<String, T> certs)
        throws GeneralSecurityException, IOException {
      KeyStore ks = createEmptyKeyStore();
      for (Map.Entry<String, T> cert : certs.entrySet()) {
        ks.setCertificateEntry(cert.getKey(), cert.getValue());
      }
      saveKeyStore(ks, trustStore, password);
    }
  
    /**
     * Create a self-signed X.509 Certificate.
     *
     * @param dn the X.509 Distinguished Name, eg "CN=Test, L=London, C=GB"
     * @param pair the KeyPair
     * @param days how many days from now the Certificate is valid for
     * @param algorithm the signing algorithm, eg "SHA1withRSA"
     * @return the self-signed certificate
     */
    public static X509Certificate generateCertificate(
        String dn, KeyPair pair, int days, String algorithm) throws Exception {
      return generateCertificate(dn, pair, days, algorithm, false, null, null, null);
    }
  
    /**
     * Create a self-signed X.509 Certificate.
     *
     * @param dn the X.509 Distinguished Name, eg "CN=Test, L=London, C=GB"
     * @param pair the KeyPair for the server
     * @param days how many days from now the Certificate is valid for
     * @param algorithm the signing algorithm, eg "SHA1withRSA"
     * @param generateCaCert Is this request to generate a CA cert
     * @param altNames Optional: Alternate names to be added to the cert - we add them as both
     *     hostnames and ip's.
     * @param caKeyPair Optional: the KeyPair of the CA, to be used to sign this certificate. caCert
     *     should also be specified to use it
     * @param caCert Optional: the CA cert, to be used to sign this certificate. caKeyPair should also
     *     be specified to use it
     * @return the signed certificate (signed using ca if provided, else self-signed)
     */
    @SuppressWarnings("JavaUtilDate")
    public static X509Certificate generateCertificate(
        String dn,
        KeyPair pair,
        int days,
        String algorithm,
        boolean generateCaCert,
        String[] altNames,
        KeyPair caKeyPair,
        X509Certificate caCert)
        throws Exception {
  
      Date from = new Date();
      Date to = new Date(from.getTime() + days * 86400000L);
      BigInteger sn = new BigInteger(64, new SecureRandom());
      X500Name subjectName = new X500Name(dn);
  
      X500Name issuerName;
      KeyPair signingKeyPair;
  
      if (caKeyPair != null && caCert != null) {
        issuerName = new JcaX509CertificateHolder(caCert).getSubject();
        signingKeyPair = caKeyPair;
      } else {
        issuerName = subjectName;
        // self signed
        signingKeyPair = pair;
      }
  
      X509v3CertificateBuilder certBuilder =
          new JcaX509v3CertificateBuilder(
              issuerName, sn, from, to, new X500Name(dn), pair.getPublic());
  
      if (null != altNames) {
        Stream<GeneralName> dnsStream =
            Arrays.stream(altNames).map(h -> new GeneralName(GeneralName.dNSName, h));
        Stream<GeneralName> ipStream =
            Arrays.stream(altNames)
                .map(
                    h -> {
                      try {
                        return new GeneralName(GeneralName.iPAddress, h);
                      } catch (Exception ex) {
                        return null;
                      }
                    })
                .filter(Objects::nonNull);
  
        GeneralName[] arr = Stream.concat(dnsStream, ipStream).toArray(GeneralName[]::new);
        GeneralNames names = new GeneralNames(arr);
  
        certBuilder.addExtension(Extension.subjectAlternativeName, false, names);
        LOG.info("Added subjectAlternativeName extension for hosts : " + Arrays.toString(altNames));
      }
  
      if (generateCaCert) {
        certBuilder.addExtension(Extension.basicConstraints, true, new BasicConstraints(true));
        LOG.info("Added CA cert extension");
      }
  
      ContentSigner signer =
          new JcaContentSignerBuilder(algorithm).build(signingKeyPair.getPrivate());
      return new JcaX509CertificateConverter().getCertificate(certBuilder.build(signer));
    }
  
    public static KeyPair generateKeyPair(String algorithm) throws NoSuchAlgorithmException {
      KeyPairGenerator keyGen = KeyPairGenerator.getInstance(algorithm);
      keyGen.initialize(4096);
      return keyGen.genKeyPair();
    }
  
    /**
     * Creates a keystore with a single key and saves it to a file.
     *
     * @param keyStore File keystore to save
     * @param password String store password to set on keystore
     * @param keyPassword String key password to set on key
     * @param alias String alias to use for the key
     * @param privateKey Key to save in keystore
     * @param cert Certificate to use as certificate chain associated to key
     * @throws GeneralSecurityException for any error with the security APIs
     * @throws IOException if there is an I/O error saving the file
     */
    public static void createKeyStore(
        File keyStore,
        String password,
        String keyPassword,
        String alias,
        Key privateKey,
        Certificate cert)
        throws GeneralSecurityException, IOException {
      KeyStore ks = createEmptyKeyStore();
      ks.setKeyEntry(alias, privateKey, keyPassword.toCharArray(), new Certificate[] {cert});
      saveKeyStore(ks, keyStore, password);
    }
  
    public static void createKeyStore(
        File keyStore, String password, String alias, Key privateKey, Certificate cert)
        throws GeneralSecurityException, IOException {
      KeyStore ks = createEmptyKeyStore();
      ks.setKeyEntry(alias, privateKey, password.toCharArray(), new Certificate[] {cert});
      saveKeyStore(ks, keyStore, password);
    }
  
    private static KeyStore createEmptyKeyStore() throws GeneralSecurityException, IOException {
      KeyStore ks = KeyStore.getInstance("PKCS12");
      ks.load(null, null); // initialize
      return ks;
    }
  
    private static void saveKeyStore(KeyStore ks, File keyStore, String password)
        throws GeneralSecurityException, IOException {
      // Write the file atomically to ensure tests don't read a partial write
      File tempFile = File.createTempFile("temp-key-store", "jks");
      FileOutputStream out = new FileOutputStream(tempFile);
      try {
        ks.store(out, password.toCharArray());
        out.close();
        Files.move(
            tempFile.toPath(),
            keyStore.toPath(),
            StandardCopyOption.REPLACE_EXISTING,
            StandardCopyOption.ATOMIC_MOVE);
      } finally {
        out.close();
      }
    }
  }
  