/*
 * Copyright 2018 Scytl Secure Electronic Voting SA
 *
 * All rights reserved
 *
 * See our extended copyright notice in *file 'Copyright.txt' which is part of this source code package
 */
package com.scytl.cryptolib.primitives.primes.utils;

import java.io.IOException;
import java.io.StringReader;
import java.io.StringWriter;
import java.security.KeyPair;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.cert.Certificate;
import java.security.cert.CertificateException;

import org.bouncycastle.asn1.x509.SubjectPublicKeyInfo;
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 org.bouncycastle.openssl.jcajce.JcaPEMWriter;
import org.bouncycastle.util.encoders.DecoderException;

import com.scytl.cryptolib.api.exceptions.GeneralCryptoLibException;
import com.scytl.cryptolib.commons.utils.validations.Validate;

/**
 * Provides utilities for writing cryptographic primitives to and reading them
 * from PEM format.
 */
public final class PemUtils {

    private static final String PEM_HEADER_PREFIX = "-----BEGIN";

    private static final String PEM_FOOTER_PREFIX = "-----END";

    private static final String PUBLIC_KEY_PEM_HEADER_SUFFIX =
        "PUBLIC KEY-----";

    private static final String PRIVATE_KEY_PEM_HEADER_SUFFIX =
        "PRIVATE KEY-----";

    private static final String CERTIFICATE_PEM_HEADER_SUFFIX =
        "CERTIFICATE-----";

    private PemUtils() {
    }

    /**
     * Converts a {@link java.security.PublicKey} to PEM format.
     *
     * @param key
     *            the {@link java.security.PublicKey} to convert.
     * @return the {@link java.security.PublicKey} in PEM format.
     * @throws GeneralCryptoLibException
     *             if the input validation or the object to PEM conversion
     *             fails.
     */
    public static final String publicKeyToPem(final PublicKey key)
            throws GeneralCryptoLibException {

        Validate.notNull(key, "Public key");
        Validate.notNullOrEmpty(key.getEncoded(), "Public key content");

        return toPem(key, PublicKey.class);
    }

    /**
     * Converts a {@link java.security.PrivateKey} to PEM format.
     *
     * @param key
     *            the {@link java.security.PrivateKey} to convert.
     * @return the {@link java.security.PrivateKey} in PEM format.
     * @throws GeneralCryptoLibException
     *             if the input validation or the object to PEM conversion
     *             fails.
     */
    public static final String privateKeyToPem(final PrivateKey key)
            throws GeneralCryptoLibException {

        Validate.notNull(key, "Private key");
        Validate.notNullOrEmpty(key.getEncoded(), "Private key content");

        return toPem(key, PrivateKey.class);
    }

    /**
     * Converts a {@link java.security.cert.Certificate} to PEM format.
     *
     * @param certificate
     *            the {@link java.security.cert.Certificate} to convert.
     * @return the {@link java.security.cert.Certificate} in PEM format.
     * @throws GeneralCryptoLibException
     *             if the input validation or the object to PEM conversion
     *             fails.
     */
    public static final String certificateToPem(
            final Certificate certificate)
            throws GeneralCryptoLibException {

        Validate.notNull(certificate, "Certificate");

        return toPem(certificate, Certificate.class);
    }

    /**
     * Retrieves a {@link java.security.PublicKey} from a string in PEM format.
     *
     * @param pemStr
     *            the {@link java.security.PublicKey} as a string in PEM format.
     * @return the {@link java.security.PublicKey}
     * @throws GeneralCryptoLibException
     *             if the input validation or the PEM to object conversion
     *             fails.
     */
    public static PublicKey publicKeyFromPem(final String pemStr)
            throws GeneralCryptoLibException {

        Validate.notNullOrBlank(pemStr, "Public key PEM string");
        validatePemFormat(pemStr, PUBLIC_KEY_PEM_HEADER_SUFFIX,
            "Public key PEM string");

        SubjectPublicKeyInfo publicKeyInfo =
            fromPem(formatPemString(pemStr, PUBLIC_KEY_PEM_HEADER_SUFFIX),
                SubjectPublicKeyInfo.class);

        try {
            return new JcaPEMKeyConverter()
                .setProvider(BouncyCastleProvider.PROVIDER_NAME)
                .getPublicKey(publicKeyInfo);
        } catch (IOException e) {
            throw new GeneralCryptoLibException(
                "Could not retrieve public key from object of type {@link "
                    + SubjectPublicKeyInfo.class + "}",
                e);
        }
    }

    /**
     * Retrieves a {@link java.security.PrivateKey} from a string in PEM format.
     *
     * @param pemStr
     *            the {@link java.security.PrivateKey} as a string in PEM
     *            format.
     * @return the {@link java.security.PrivateKey}
     * @throws GeneralCryptoLibException
     *             if the input validation or the PEM to object conversion
     *             fails.
     */
    public static PrivateKey privateKeyFromPem(final String pemStr)
            throws GeneralCryptoLibException {

        Validate.notNullOrBlank(pemStr, "Private key PEM string");
        validatePemFormat(pemStr, PRIVATE_KEY_PEM_HEADER_SUFFIX,
            "Private key PEM string");

        PEMKeyPair pemKeyPair =
            fromPem(formatPemString(pemStr, PRIVATE_KEY_PEM_HEADER_SUFFIX),
                PEMKeyPair.class);

        try {
            KeyPair keyPair = new JcaPEMKeyConverter()
                .setProvider(BouncyCastleProvider.PROVIDER_NAME)
                .getKeyPair(pemKeyPair);

            return keyPair.getPrivate();
        } catch (IOException e) {
            throw new GeneralCryptoLibException(
                "Could not retrieve private key from object of type {@link "
                    + PEMKeyPair.class + "}",
                e);
        }
    }

    /**
     * Retrieves a {java.security.cert.Certificate} from a string in PEM format.
     *
     * @param pemStr
     *            the {java.security.cert.Certificate} as a string in PEM
     *            format.
     * @return the {java.security.cert.Certificate}
     * @throws GeneralCryptoLibException
     *             if the input validation or the PEM to object conversion
     *             fails.
     */
    public static Certificate certificateFromPem(final String pemStr)
            throws GeneralCryptoLibException {

        Validate.notNullOrBlank(pemStr, "Certificate PEM string");
        validatePemFormat(pemStr, CERTIFICATE_PEM_HEADER_SUFFIX,
            "Certificate PEM string");

        X509CertificateHolder certificateHolder =
            fromPem(formatPemString(pemStr, CERTIFICATE_PEM_HEADER_SUFFIX),
                X509CertificateHolder.class);

        try {
            return new JcaX509CertificateConverter()
                .getCertificate(certificateHolder);
        } catch (CertificateException e) {
            throw new GeneralCryptoLibException(
                "Could not retrieve public key from object of type {@link "
                    + X509CertificateHolder.class + "}",
                e);
        }
    }

    private static <T> String toPem(final Object object,
            final Class<T> objectClass) throws GeneralCryptoLibException {

        StringWriter pemStringWriter = new StringWriter();
        try (JcaPEMWriter writer = new JcaPEMWriter(pemStringWriter)) {
            writer.writeObject(object);
        } catch (IOException e) {
            throw new GeneralCryptoLibException(
                "Could not convert object of type {@link " + objectClass
                    + "} to PEM format.",
                e);
        }

        return pemStringWriter.toString();
    }

    private static <T> T fromPem(final String pemStr,
            final Class<T> objectClass) throws GeneralCryptoLibException {

        try (PEMParser parser = new PEMParser(new StringReader(pemStr))) {
            Object object = parser.readObject();

            return objectClass.cast(object);
        } catch (IOException | DecoderException e) {
            throw new GeneralCryptoLibException(
                "Could not convert PEM string " + pemStr
                    + " to object of type {@link " + objectClass + "}",
                e);
        }
    }

    private static final void validatePemFormat(final String pemStr,
            final String headerSuffix, final String label)
            throws GeneralCryptoLibException {

        if (!pemStr.contains(PEM_HEADER_PREFIX)
            || !pemStr.contains(headerSuffix)) {
            throw new GeneralCryptoLibException(
                label + " does not contain valid header.");
        }

        if (!pemStr.contains(PEM_FOOTER_PREFIX)) {
            throw new GeneralCryptoLibException(
                label + " does not contain valid footer.");
        }
    }

    private static String formatPemString(final String pemStr,
            final String headerSuffix) {

        String formattedPemStr =
            pemStr.replace("\n", "").replace("\r", "");

        formattedPemStr = formattedPemStr.replaceFirst(headerSuffix,
            headerSuffix + "\n");

        formattedPemStr = formattedPemStr.replaceFirst(PEM_FOOTER_PREFIX,
            "\n" + PEM_FOOTER_PREFIX);

        return formattedPemStr;
    }
}
