package com.zy.integration.iot.util;

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.PEMKeyPair;
import org.bouncycastle.openssl.PEMParser;
import org.bouncycastle.openssl.jcajce.JcaPEMKeyConverter;

import javax.net.ssl.KeyManagerFactory;
import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLSocketFactory;
import javax.net.ssl.TrustManagerFactory;
import java.io.BufferedReader;
import java.io.FileReader;
import java.security.KeyStore;
import java.security.PrivateKey;
import java.security.SecureRandom;
import java.security.Security;
import java.security.cert.Certificate;
import java.security.cert.X509Certificate;

public final class IotMqttSslUtils {

    private static final String BC_PROVIDER = "BC";

    private IotMqttSslUtils() {
    }

    public static SSLSocketFactory buildSocketFactory(String caCertPath, String clientCertPath, String privateKeyPath) throws Exception {
        ensureProvider();
        X509Certificate caCertificate = readCertificate(caCertPath);
        X509Certificate clientCertificate = readCertificate(clientCertPath);
        PrivateKey privateKey = readPrivateKey(privateKeyPath);

        KeyStore clientStore = KeyStore.getInstance(KeyStore.getDefaultType());
        clientStore.load(null, null);
        clientStore.setKeyEntry("iot-client-key", privateKey, new char[0], new Certificate[]{clientCertificate});

        KeyManagerFactory keyManagerFactory = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm());
        keyManagerFactory.init(clientStore, new char[0]);

        KeyStore trustStore = KeyStore.getInstance(KeyStore.getDefaultType());
        trustStore.load(null, null);
        trustStore.setCertificateEntry("iot-ca", caCertificate);

        TrustManagerFactory trustManagerFactory = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
        trustManagerFactory.init(trustStore);

        SSLContext sslContext = SSLContext.getInstance("TLSv1.2");
        sslContext.init(keyManagerFactory.getKeyManagers(), trustManagerFactory.getTrustManagers(), new SecureRandom());
        return sslContext.getSocketFactory();
    }

    private static X509Certificate readCertificate(String path) throws Exception {
        try (PEMParser parser = new PEMParser(new BufferedReader(new FileReader(path)))) {
            Object object = parser.readObject();
            if (!(object instanceof X509CertificateHolder)) {
                throw new IllegalArgumentException("invalid certificate file: " + path);
            }
            return new JcaX509CertificateConverter()
                    .setProvider(BC_PROVIDER)
                    .getCertificate((X509CertificateHolder) object);
        }
    }

    private static PrivateKey readPrivateKey(String path) throws Exception {
        try (PEMParser parser = new PEMParser(new BufferedReader(new FileReader(path)))) {
            Object object = parser.readObject();
            JcaPEMKeyConverter converter = new JcaPEMKeyConverter().setProvider(BC_PROVIDER);
            if (object instanceof PEMKeyPair) {
                return converter.getKeyPair((PEMKeyPair) object).getPrivate();
            }
            if (object instanceof PrivateKeyInfo) {
                return converter.getPrivateKey((PrivateKeyInfo) object);
            }
            throw new IllegalArgumentException("invalid private key file: " + path);
        }
    }

    private static void ensureProvider() {
        if (Security.getProvider(BC_PROVIDER) == null) {
            Security.addProvider(new BouncyCastleProvider());
        }
    }
}