package org.wso2.carbon.device.mgt.ios.core.impl;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.bouncycastle.cert.jcajce.JcaCertStore;
import org.bouncycastle.cms.CMSAlgorithm;
import org.bouncycastle.cms.CMSEnvelopedData;
import org.bouncycastle.cms.CMSEnvelopedDataGenerator;
import org.bouncycastle.cms.CMSException;
import org.bouncycastle.cms.CMSProcessableByteArray;
import org.bouncycastle.cms.CMSSignedData;
import org.bouncycastle.cms.CMSSignedDataGenerator;
import org.bouncycastle.cms.CMSTypedData;
import org.bouncycastle.cms.jcajce.JcaSignerInfoGeneratorBuilder;
import org.bouncycastle.cms.jcajce.JceCMSContentEncryptorBuilder;
import org.bouncycastle.cms.jcajce.JceKeyTransRecipientInfoGenerator;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.operator.ContentSigner;
import org.bouncycastle.operator.OperatorCreationException;
import org.bouncycastle.operator.jcajce.JcaContentSignerBuilder;
import org.bouncycastle.operator.jcajce.JcaDigestCalculatorProviderBuilder;
import org.bouncycastle.util.Store;
import org.wso2.carbon.device.mgt.ios.core.exception.IOSEnrollmentException;
import org.wso2.carbon.device.mgt.ios.core.util.AppConfigurations;

import java.io.IOException;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.Security;
import java.security.UnrecoverableKeyException;
import java.security.cert.Certificate;
import java.security.cert.CertificateEncodingException;
import java.security.cert.X509Certificate;
import java.util.ArrayList;
import java.util.List;

public class PKCSSigner {

    private static final Log log = LogFactory.getLog(PKCSSigner.class);

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

	public CMSSignedDataGenerator setUpProvider(final KeyStore keystore)
			throws IOSEnrollmentException {

		CMSSignedDataGenerator generator;
		try {
			Certificate[] certificateChain = keystore.getCertificateChain(
                    AppConfigurations.getConfigEntry(AppConfigurations.KEY_ALIAS_KEYSTORE));

			final List<Certificate> certificateList = new ArrayList<Certificate>();

			for (int i = 0, length = certificateChain == null ? 0 : certificateChain.length; i < length; i++) {
				certificateList.add(certificateChain[i]);
			}

			Store store = new JcaCertStore(certificateList);
			Certificate certificate = keystore.getCertificate(AppConfigurations
					.getConfigEntry(AppConfigurations.KEY_ALIAS_KEYSTORE));

			ContentSigner contentSigner = new JcaContentSignerBuilder(
					AppConfigurations.SIGNATUREALGO)
					.setProvider(AppConfigurations.PROVIDER)
					.build((PrivateKey) (keystore.getKey(
							AppConfigurations.getConfigEntry(AppConfigurations.KEY_ALIAS_KEYSTORE),
							AppConfigurations.getConfigEntry(AppConfigurations.KEYSTORE_KEY_PASSWORD).toCharArray())));
			generator = new CMSSignedDataGenerator();
			generator.addSignerInfoGenerator(new JcaSignerInfoGeneratorBuilder(
					new JcaDigestCalculatorProviderBuilder().setProvider(
							AppConfigurations.PROVIDER).build()).build(contentSigner,
					(X509Certificate) certificate));

			generator.addCertificates(store);
		} catch (KeyStoreException e) {
            String errorMsg = "KeyStore issue occurred in setUpProvider";
            log.error(errorMsg, e);
            throw new IOSEnrollmentException(errorMsg, e);
		} catch (CertificateEncodingException e) {
            String errorMsg = "Certificate issue issue occurred in setUpProvider";
            log.error(errorMsg, e);
            throw new IOSEnrollmentException(errorMsg, e);
		} catch (UnrecoverableKeyException e) {
            String errorMsg = "Unrecoverable key in setUpProvider";
            log.error(errorMsg, e);
            throw new IOSEnrollmentException(errorMsg, e);
		} catch (OperatorCreationException e) {
            String errorMsg = "Operator creation issue occurred in setUpProvider";
            log.error(errorMsg, e);
            throw new IOSEnrollmentException(errorMsg, e);
		} catch (NoSuchAlgorithmException e) {
            String errorMsg = "No such algorithm in setUpProvider";
            log.error(errorMsg, e);
            throw new IOSEnrollmentException(errorMsg, e);
		} catch (CMSException e) {
            String errorMsg = "CMS issue occurred in setUpProvider";
            log.error(errorMsg, e);
            throw new IOSEnrollmentException(errorMsg, e);
		}

		return generator;
	}

	public byte[] signPKCS(final byte[] content, final CMSSignedDataGenerator generator) throws IOSEnrollmentException {

		try {
			CMSTypedData cmsData = new CMSProcessableByteArray(content);
			CMSSignedData signedData = generator.generate(cmsData, true);

			return signedData.getEncoded();
		} catch (CMSException e) {
            String errorMsg = "CMS issue occurred in signPKCS";
            log.error(errorMsg, e);
            throw new IOSEnrollmentException(errorMsg, e);
		} catch (IOException e) {
            String errorMsg = "Input output issue occurred in signPKCS";
            log.error(errorMsg, e);
            throw new IOSEnrollmentException(errorMsg, e);
		}
	}

	public byte[] getSignedData(byte[] input) throws IOSEnrollmentException {
        KeyStoreReader keyStoreReader = new KeyStoreReader();
        KeyStore keyStore = keyStoreReader.loadCarbonKeyStore();
        CMSSignedDataGenerator signedDataGenerator = setUpProvider(keyStore);
        return signPKCS(input, signedDataGenerator);
	}

	public byte[] getEncryptedData(byte[] input, X509Certificate certificate) throws IOSEnrollmentException {

        try {
            CMSTypedData msg = new CMSProcessableByteArray(input);

            CMSEnvelopedDataGenerator edGen = new CMSEnvelopedDataGenerator();
            edGen.addRecipientInfoGenerator(new JceKeyTransRecipientInfoGenerator(
                    certificate).setProvider(AppConfigurations.PROVIDER));
            CMSEnvelopedData envelopedData = edGen.generate(msg,
                    new JceCMSContentEncryptorBuilder(
                            CMSAlgorithm.DES_EDE3_CBC).setProvider(AppConfigurations.PROVIDER).build());

            return envelopedData.getEncoded();
        } catch (CertificateEncodingException e) {
            String errorMsg = "Certificate encoding issue when generating operation payload";
            log.error(errorMsg, e);
            throw new IOSEnrollmentException(errorMsg, e);
        } catch (CMSException e) {
            String errorMsg = "CMSException when generating operation payload";
            log.error(errorMsg, e);
            throw new IOSEnrollmentException(errorMsg, e);
        } catch (IOException e) {
            String errorMsg = "IOException when generating operation payload";
            log.error(errorMsg, e);
            throw new IOSEnrollmentException(errorMsg, e);
        }
	}
}
