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

import org.apache.commons.codec.binary.Base64;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.bouncycastle.asn1.x500.X500Name;
import org.bouncycastle.cert.X509CertificateHolder;
import org.bouncycastle.cert.jcajce.JcaX509CertificateConverter;
import org.bouncycastle.cert.jcajce.JcaX509CertificateHolder;
import org.bouncycastle.cms.CMSAlgorithm;
import org.bouncycastle.cms.CMSEnvelopedData;
import org.bouncycastle.cms.CMSEnvelopedDataGenerator;
import org.bouncycastle.cms.CMSException;
import org.bouncycastle.cms.CMSProcessable;
import org.bouncycastle.cms.CMSProcessableByteArray;
import org.bouncycastle.cms.CMSSignedData;
import org.bouncycastle.cms.CMSSignedDataGenerator;
import org.bouncycastle.cms.CMSTypedData;
import org.bouncycastle.cms.SignerInformation;
import org.bouncycastle.cms.SignerInformationStore;
import org.bouncycastle.cms.jcajce.JcaSimpleSignerInfoVerifierBuilder;
import org.bouncycastle.cms.jcajce.JceCMSContentEncryptorBuilder;
import org.bouncycastle.cms.jcajce.JceKeyTransRecipientInfoGenerator;
import org.bouncycastle.operator.OperatorCreationException;
import org.bouncycastle.util.Store;
import org.wso2.carbon.certificate.mgt.core.exception.KeystoreException;
import org.wso2.carbon.device.mgt.common.DeviceManagementException;
import org.wso2.carbon.device.mgt.ios.core.exception.IOSEnrollmentException;
import org.wso2.carbon.device.mgt.ios.core.internal.IOSEnrollmentServiceHolder;
import org.wso2.carbon.device.mgt.ios.core.publisher.TokenPersistence;
import org.wso2.carbon.device.mgt.ios.core.util.AppConfigurations;
import org.wso2.carbon.device.mgt.ios.payload.dto.CertificateAttributes;
import org.wso2.carbon.device.mgt.ios.payload.dto.DeviceProperties;
import org.wso2.carbon.device.mgt.ios.payload.enrollment.EnrollmentPayloadGenerator;
import org.wso2.carbon.device.mgt.ios.payload.enrollment.EnrollmentAttributeExtractor;
import org.wso2.carbon.device.mgt.ios.payload.exception.PListException;

import java.io.IOException;
import java.io.InputStream;
import java.io.UnsupportedEncodingException;
import java.security.KeyStore;
import java.security.cert.CertificateEncodingException;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;
import java.util.Collection;
import java.util.Iterator;

public class RequestHandler {

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

    /**
     * Get the serial number from the profile request sent to the server by the device.
     * @param inputStream the inputstream recieved from the device request.
     * @return the serial number fetched from the payload.
     * @throws IOSEnrollmentException
     */
    public DeviceProperties getDeviceProperties(InputStream inputStream) throws IOSEnrollmentException {
        DeviceProperties deviceDetails = null;
        try {
            CMSSignedData cms = new CMSSignedData(inputStream);
            Store store = cms.getCertificates();
            SignerInformationStore signers = cms.getSignerInfos();

            if (signers == null) {
                throw new IOSEnrollmentException("Signer information store is null in " +
                        "handleProfileRequest");
            }

            if (store == null) {
                throw new IOSEnrollmentException("Store is null in handleProfileRequest");
            }

            Collection<?> collectionSigners = signers.getSigners();
            if (collectionSigners == null) {
                throw new IOSEnrollmentException("Signer collection is null in " +
                        "handleProfileRequest");
            }

            for (SignerInformation signer : (Collection<SignerInformation>) collectionSigners) {

                Collection<?> certCollection = store.getMatches(signer.getSID());
                if (certCollection == null) {
                    throw new IOSEnrollmentException("Certificate collection is null in " +
                            "handleProfileRequest");
                }

                Iterator<?> certIt = certCollection.iterator();
                X509CertificateHolder certHolder = (X509CertificateHolder) certIt.next();
                X509Certificate certificate = new JcaX509CertificateConverter()
                        .setProvider(AppConfigurations.PROVIDER).getCertificate(certHolder);

                // Verifies the signer
                if (signer.verify(new JcaSimpleSignerInfoVerifierBuilder()
                        .setProvider(AppConfigurations.PROVIDER).build(certificate))) {
                    CMSProcessable cmsProcessable = cms.getSignedContent();
                    String contentString = new String((byte[]) cmsProcessable.getContent(),
                            AppConfigurations.UTF_8);
                    EnrollmentAttributeExtractor enrollmentAttributeExtractor = new
                            EnrollmentAttributeExtractor();
                    deviceDetails = enrollmentAttributeExtractor.extractDeviceProperties(contentString);
                }
            }
        } catch (CMSException e) {
            String errorMsg = "CMS issue occurred when handling profile request";
            log.error(errorMsg, e);
            throw new IOSEnrollmentException(errorMsg, e);
        } catch (PListException e) {
            String errorMsg = "PList generation issue occurred when handling profile request";
            log.error(errorMsg, e);
            throw new IOSEnrollmentException(errorMsg, e);
        } catch (OperatorCreationException e) {
            String errorMsg = "Operator creation issue occurred when handling profile request";
            log.error(errorMsg, e);
            throw new IOSEnrollmentException(errorMsg, e);
        } catch (CertificateException e) {
            String errorMsg = "Certificate issue occurred when handling profile request";
            log.error(errorMsg, e);
            throw new IOSEnrollmentException(errorMsg, e);
        } catch (UnsupportedEncodingException e) {
            String errorMsg = "Unsupported encoding in handleProfileRequest";
            log.error(errorMsg, e);
            throw new IOSEnrollmentException(errorMsg, e);
        }
        return deviceDetails;
    }

    public byte[] handleProfileRequest(InputStream inputStream, String tenant, String topicId,
                                       CertificateAttributes certAttributes, boolean isDEPRequest)
            throws IOSEnrollmentException {

        try {
            CMSSignedData cms = new CMSSignedData(inputStream);
            Store store = cms.getCertificates();
            SignerInformationStore signers = cms.getSignerInfos();

            if (signers == null) {
                throw new IOSEnrollmentException("Signer information store is null in handleProfileRequest");
            }

            if (store == null) {
                throw new IOSEnrollmentException("Store is null in handleProfileRequest");
            }

            Collection<?> collectionSigners = signers.getSigners();
            if (collectionSigners == null) {
                throw new IOSEnrollmentException("Signer collection is null in handleProfileRequest");
            }

            for (SignerInformation signer : (Collection<SignerInformation>) collectionSigners) {

                Collection<?> certCollection = store.getMatches(signer.getSID());
                if (certCollection == null) {
                    throw new IOSEnrollmentException("Certificate collection is null in handleProfileRequest");
                }

                Iterator<?> certIt = certCollection.iterator();
                X509CertificateHolder certHolder = (X509CertificateHolder) certIt.next();
                X509Certificate certificate = new JcaX509CertificateConverter().setProvider(AppConfigurations.PROVIDER)
                        .getCertificate(certHolder);

                // verifies the signer
                if (signer.verify(new JcaSimpleSignerInfoVerifierBuilder().setProvider(AppConfigurations.PROVIDER)
                        .build(certificate))) {

                    KeyStoreReader keyStoreReader = new KeyStoreReader();
                    X509Certificate caCertificate = (X509Certificate) IOSEnrollmentServiceHolder.getInstance().
                            getCertificateManagementService().getCACertificate();

                    JcaX509CertificateHolder jcaX509CertificateHolder = new JcaX509CertificateHolder(caCertificate);
                    X500Name x500name = jcaX509CertificateHolder.getSubject();

                    EnrollmentPayloadGenerator enrollmentPayloadGenerator = new EnrollmentPayloadGenerator();
                    PKCSSigner pkcsSigner = new PKCSSigner();
                    KeyStore keyStore = keyStoreReader.loadCarbonKeyStore();
                    CMSSignedDataGenerator signedDataGenerator = pkcsSigner.setUpProvider(keyStore);

                    // extract content
                    CMSProcessable cmsProcessable = cms.getSignedContent();
                    String contentString;

                    try {
                        contentString = new String((byte[]) cmsProcessable.getContent(), AppConfigurations.UTF_8);
                    } catch (UnsupportedEncodingException e) {
                        throw new IOSEnrollmentException("Unsupported encoding in handleProfileRequest");
                    }

                    EnrollmentAttributeExtractor enrollmentAttributeExtractor = new EnrollmentAttributeExtractor();
                    DeviceProperties deviceProperties = enrollmentAttributeExtractor.
                            extractDeviceProperties(contentString);

                    byte[] signedData;
                    if (isDEPRequest || signer.getSID().getIssuer().equals(x500name)) {
                        signedData = signPKCSMDMConfigurationPayload(pkcsSigner, enrollmentPayloadGenerator,
                                deviceProperties, tenant, topicId, certAttributes, certificate, signedDataGenerator);
                        //Device enrollment details needs to be added in DEP in the first profile
                        // delivery, as in DEP the final enrollment profile is delivered in the first
                        // call to profile endpoint.
                        if (isDEPRequest) {
                            TokenPersistence tokenPersistence = new TokenPersistence();
                            tokenPersistence.saveDeviceTokens(deviceProperties);
                        }
                    } else {
                        signedData = signPKCSWithTokenPublish(pkcsSigner, enrollmentPayloadGenerator, deviceProperties,
                                tenant, certAttributes, signedDataGenerator);
                    }

                    return signedData;
                }
            }

        } catch (OperatorCreationException e) {
            String errorMsg = "Operator creation issue occurred when handling profile request";
            log.error(errorMsg, e);
            throw new IOSEnrollmentException(errorMsg, e);
        } catch (CertificateException e) {
            String errorMsg = "Certificate issue occurred when handling profile request";
            log.error(errorMsg, e);
            throw new IOSEnrollmentException(errorMsg, e);
        } catch (KeystoreException e) {
            String errorMsg = "Keystore reading error occurred when handling profile request";
            log.error(errorMsg, e);
            throw new IOSEnrollmentException(errorMsg, e);
        } catch (CMSException e) {
            String errorMsg = "CMS issue occurred when handling profile request";
            log.error(errorMsg, e);
            throw new IOSEnrollmentException(errorMsg, e);
        } catch (PListException e) {
            String errorMsg = "PList generation issue occurred when handling profile request";
            log.error(errorMsg, e);
            throw new IOSEnrollmentException(errorMsg, e);
        } catch (DeviceManagementException e) {
            String errorMsg = "Error while saving the tokens.";
            log.error(errorMsg, e);
            throw new IOSEnrollmentException(errorMsg, e);
        }

        return null;
    }



    private byte[] signPKCSMDMConfigurationPayload(PKCSSigner pkcsSigner,
                                                   EnrollmentPayloadGenerator enrollmentPayloadGenerator,
                                                   DeviceProperties deviceProperties, String tenant, String topicId,
                                                   CertificateAttributes certAttributes, X509Certificate certificate,
                                                   CMSSignedDataGenerator signedDataGenerator)
            throws IOSEnrollmentException, PListException, CMSException, CertificateEncodingException {
        KeyStoreReader keyStoreReader = new KeyStoreReader();

        String mdmConfigurationPayload = enrollmentPayloadGenerator
                .generateMDMConfigurationPayload(
                        deviceProperties.getDeviceIdentifier(), tenant,
                        topicId, certAttributes,
                        AppConfigurations.getConfigEntry(AppConfigurations.CHECKIN_URL),
                        AppConfigurations.getConfigEntry(AppConfigurations.SERVER_URL),
                        AppConfigurations.getConfigEntry(AppConfigurations.ENROLL_URL),
                        keyStoreReader.getEncodedCACert()
                );

        CMSTypedData msg;
        try {
            msg = new CMSProcessableByteArray(mdmConfigurationPayload.getBytes(AppConfigurations.UTF_8));
        } catch (UnsupportedEncodingException e) {
            String errorMsg = "Unsupported encoding when generating mdm configuration payload";
            log.error(errorMsg, e);
            throw new IOSEnrollmentException(errorMsg, e);
        }

        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());

        String encryptedConfigurationPayload;
        try {
            encryptedConfigurationPayload = enrollmentPayloadGenerator.generateEncryptedConfigurationPayload(
                    new String(Base64.encodeBase64(envelopedData.getEncoded()), AppConfigurations.UTF_8),
                    tenant);
        } catch (UnsupportedEncodingException e) {
            String errorMsg = "Unsupported encoding when generating encrypted configuration payload";
            log.error(errorMsg, e);
            throw new IOSEnrollmentException(errorMsg, e);
        } catch (IOException e) {
            String errorMsg = "Input output issue occurred when generating encrypted configuration payload";
            log.error(errorMsg, e);
            throw new IOSEnrollmentException(errorMsg, e);
        } catch (PListException e) {
            String errorMsg = "PList generation issue occurred when generating encrypted configuration payload";
            log.error(errorMsg, e);
            throw new IOSEnrollmentException(errorMsg, e);
        }

        return pkcsSigner.signPKCS(encryptedConfigurationPayload.getBytes(), signedDataGenerator);
    }

    private byte[] signPKCSWithTokenPublish(PKCSSigner pkcsSigner,
            EnrollmentPayloadGenerator enrollmentPayloadGenerator, DeviceProperties deviceProperties, String tenant,
            CertificateAttributes certAttributes, CMSSignedDataGenerator signedDataGenerator)
            throws IOSEnrollmentException {

        try {
            String encryptionCertificatePayload = enrollmentPayloadGenerator
                    .generateEncryptedCertificatePayload(deviceProperties.getChallenge(), tenant, certAttributes,
                            AppConfigurations.getConfigEntry(AppConfigurations.ENROLL_URL));

            byte[] signedData = pkcsSigner.signPKCS(encryptionCertificatePayload.getBytes(), signedDataGenerator);
            TokenPersistence tokenPersistence = new TokenPersistence();
            tokenPersistence.saveDeviceTokens(deviceProperties);

            return signedData;
        } catch (PListException e) {
            String errorMsg = "Error in generating encryption certificate payload";
            log.error(errorMsg, e);
            throw new IOSEnrollmentException(errorMsg, e);
        } catch (DeviceManagementException e) {
            String errorMsg = "Error occurred while saving device information";
            log.error(errorMsg, e);
            throw new IOSEnrollmentException(errorMsg, e);
        }
    }

	public String getChallengeTokenFromProfile(InputStream inputStream)
			throws PListException {

		try {
			CMSSignedData cms = new CMSSignedData(inputStream);
			CMSProcessable cmsProcessable = cms.getSignedContent();
			String contentString = new String((byte[]) cmsProcessable.getContent(), AppConfigurations.UTF_8);

			EnrollmentAttributeExtractor enrollmentAttributeExtractor = new EnrollmentAttributeExtractor();
			DeviceProperties deviceProperties = enrollmentAttributeExtractor.extractDeviceProperties(contentString);

            return deviceProperties.getChallenge();
		} catch (CMSException e) {
            String errorMsg = "CMS issue occurred in get getChallengeTokenFromProfile";
            log.error(errorMsg, e);
            throw new PListException(errorMsg, e);
		} catch (UnsupportedEncodingException e) {
            String errorMsg = "Unsupported encoding issue occurred in getChallengeTokenFromProfile";
            log.error(errorMsg, e);
            throw new PListException(errorMsg, e);
		}
	}

    public DeviceProperties getDeviceIdentifierFromProfile(InputStream inputStream) throws PListException {
        try {
            CMSSignedData cms = new CMSSignedData(inputStream);
            CMSProcessable cmsProcessable = cms.getSignedContent();
            String contentString = new String((byte[]) cmsProcessable.getContent(), AppConfigurations.UTF_8);
            EnrollmentAttributeExtractor enrollmentAttributeExtractor = new EnrollmentAttributeExtractor();
            return enrollmentAttributeExtractor.extractDeviceProperties(contentString);
        } catch (CMSException e) {
            String errorMsg = "CMS issue occurred in get getDeviceIdentifierFromProfile";
            log.error(errorMsg, e);
            throw new PListException(errorMsg, e);
        } catch (UnsupportedEncodingException e) {
            String errorMsg = "Unsupported encoding issue occurred in getDeviceIdentifierFromProfile";
            log.error(errorMsg, e);
            throw new PListException(errorMsg, e);
        }
    }
}
