package be.fgov.ehealth.technicalconnector.jca;

import be.ehealth.technicalconnector.service.sts.security.Credential;
import be.ehealth.technicalconnector.service.sts.security.KeyStoreInfo;
import be.ehealth.technicalconnector.service.sts.security.impl.BeIDCredential;
import be.ehealth.technicalconnector.service.sts.security.impl.KeyStoreCredential;
import be.ehealth.technicalconnector.session.Session;
import be.ehealth.technicalconnector.session.SessionManager;
import be.fgov.ehealth.technicalconnector.distributedkeys.DistributedKeyStoreRegistry;
import be.fgov.ehealth.technicalconnector.distributedkeys.proxy.CredentialProxy;
import be.fgov.ehealth.technicalconnector.jca.proxy.BeIDDistributedSigner;
import be.fgov.ehealth.technicalconnector.tests.junit.rule.SessionRule;
import be.fgov.ehealth.technicalconnector.tests.utils.AssumeTools;
import be.fgov.ehealth.technicalconnector.tests.utils.LoggingUtils;
import be.fgov.ehealth.technicalconnector.tests.utils.TestPropertiesLoader;
import be.fgov.ehealth.technicalconnector.tests.utils.XmlAsserter;
import org.apache.commons.lang.ArrayUtils;
import org.apache.xml.security.Init;
import org.apache.xml.security.c14n.Canonicalizer;
import org.apache.xml.security.signature.XMLSignature;
import org.apache.xml.security.transforms.Transforms;
import org.apache.xml.security.utils.Constants;
import org.apache.xml.security.utils.ElementProxy;
import org.junit.After;
import org.junit.Assert;
import org.junit.BeforeClass;
import org.junit.ClassRule;
import org.junit.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.w3c.dom.Document;

import javax.xml.parsers.DocumentBuilderFactory;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.InputStream;
import java.security.Key;
import java.security.KeyStore;
import java.security.cert.Certificate;
import java.security.cert.X509Certificate;
import java.util.Properties;

/**
 * DistributedKeysTest
 *
 * @author EHP
 */
public class DistributedKeysTest {


    private static final Logger LOG = LoggerFactory.getLogger(DistributedKeysTest.class);

    @ClassRule
    public static SessionRule rule = SessionRule.withActiveSession().build();

    @BeforeClass
    public static void init() {
        LoggingUtils.bootstrap();
    }

    private static String signFile(InputStream xmlFile, Certificate[] certificate, Key privateKey) throws Exception {
        final Document doc = DocumentBuilderFactory.newInstance()
                .newDocumentBuilder()
                .parse(xmlFile);
        Init.init();
        ElementProxy.setDefaultPrefix(Constants.SignatureSpecNS, "");
        final XMLSignature sig = new XMLSignature(doc, null, XMLSignature.ALGO_ID_SIGNATURE_RSA);
        final Transforms transforms = new Transforms(doc);
        transforms.addTransform(Transforms.TRANSFORM_ENVELOPED_SIGNATURE);
        sig.addDocument("", transforms, Constants.ALGO_ID_DIGEST_SHA1);
        for (int i = 0; i < certificate.length; i++) {
            X509Certificate cert = (X509Certificate) certificate[i];
            sig.addKeyInfo(cert);
            System.out.println(cert.getSubjectX500Principal()
                    .getName());

        }
        sig.sign(privateKey);
        doc.getDocumentElement()
                .appendChild(sig.getElement());
        final ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        outputStream.write(Canonicalizer.getInstance(Canonicalizer.ALGO_ID_C14N_WITH_COMMENTS)
                .canonicalizeSubtree(doc));
        return new String(outputStream.toByteArray());
    }

    @After
    public void tearDown() {
        DistributedKeyStoreRegistry.getInstance()
                .flushCache();

    }

    @Test
    public void signWithEid() throws Exception {
        AssumeTools.isEIDEnabled();
        Credential credential = BeIDCredential.getInstance("test", BeIDCredential.EID_AUTH_ALIAS);
        DistributedKeyStoreRegistry.getInstance()
                .createDistributedKeyStore("store-eid", new CredentialProxy(credential));
        String distributedResult = sign("Authentication", "store-eid");
        String localResult = sign(BeIDCredential.EID_AUTH_ALIAS, credential.getKeyStore(), "".toCharArray());
        XmlAsserter.assertSimilar(distributedResult, localResult);


    }

    @Test
    public void createDistributedSessionWithEid() throws Exception {

        AssumeTools.isEIDEnabled();
        DistributedKeyStoreRegistry registry = DistributedKeyStoreRegistry.getInstance();
        registry.createDistributedKeyStore(Session.IDENTIFICATION_KEYSTORE, new BeIDDistributedSigner());
        registry.createDistributedKeyStore(Session.IDENTIFICATION_KEYSTORE, new BeIDDistributedSigner());
        Assert.assertFalse(registry.getDistributedKeyStores()
                .isEmpty());
        registry.createDistributedKeyStore(Session.HOLDEROFKEY_KEYSTORE, new BeIDDistributedSigner());
        registry.createDistributedKeyStore(Session.ENCRYPTION_KEYSTORE, new BeIDDistributedSigner());

        SessionManager sessionMgr = Session.getInstance();
        sessionMgr.setKeyStore(registry.getDistributedKeyStores());
        Assert.assertFalse(registry.getDistributedKeyStores()
                .isEmpty());

        sessionMgr.createSessionEidOnly();
    }

    @Test
    public void getCertificateChainWithEid() throws Exception {
        AssumeTools.isEIDEnabled();
        Credential credential = BeIDCredential.getInstance("test", BeIDCredential.EID_AUTH_ALIAS);
        KeyStore keystore = DistributedKeyStoreRegistry.getInstance()
                .createDistributedKeyStore("store-eid", new CredentialProxy(credential));
        Certificate[] certs = keystore.getCertificateChain(BeIDCredential.EID_AUTH_ALIAS);
        Assert.assertEquals(3, certs.length);

    }

    @Test
    public void getCertificateChainWithX509() throws Exception {
        Properties props = TestPropertiesLoader.getProperties("/be.ehealth.technicalconnector.test.properties");
        String eHealthCertificate = props.getProperty("test.keystore.location");
        String eHealthCertificatePwd = props.getProperty("test.keystore.password");
        String alias = props.getProperty("test.keystore.alias");
        KeyStoreInfo ksInfo = new KeyStoreInfo(eHealthCertificate, eHealthCertificatePwd.toCharArray(), alias, eHealthCertificatePwd.toCharArray());
        Credential credential = new KeyStoreCredential(ksInfo);

        KeyStore keystore = DistributedKeyStoreRegistry.getInstance()
                .createDistributedKeyStore("store-x509", new CredentialProxy(credential));
        Certificate[] certs = keystore.getCertificateChain(alias);
        Assert.assertEquals(2, certs.length);
        for (int i = 0; i < certs.length; i++) {
            X509Certificate cert = (X509Certificate) certs[i];
            System.out.println(cert.getSubjectX500Principal()
                    .getName());
        }

    }

    @Test
    public void signWithX509() throws Exception {
        Properties props = TestPropertiesLoader.getProperties("/be.ehealth.technicalconnector.test.properties");
        String eHealthCertificate = props.getProperty("test.keystore.location");
        String eHealthCertificatePwd = props.getProperty("test.keystore.password");
        String alias = props.getProperty("test.keystore.alias");
        KeyStoreInfo ksInfo = new KeyStoreInfo(eHealthCertificate, eHealthCertificatePwd.toCharArray(), alias, eHealthCertificatePwd.toCharArray());
        Credential credential = new KeyStoreCredential(ksInfo);

        DistributedKeyStoreRegistry.getInstance()
                .createDistributedKeyStore("store-jks", new CredentialProxy(credential));
        String distributedResult = sign("authentication", "store-jks");
        String localResult = sign("authentication", credential.getKeyStore(), eHealthCertificatePwd.toCharArray());
        XmlAsserter.assertSimilar(localResult, distributedResult);

    }

    private String sign(String alias, KeyStore store, char[] password) throws Exception {
        InputStream xmlFile = new ByteArrayInputStream("<boe/>".getBytes());
        String result = signFile(xmlFile, store.getCertificateChain(alias), store.getKey(alias, password));
        LOG.debug(result);
        return result;
    }

    private String sign(String alias, String storeName) throws Exception {
        KeyStore store = DistributedKeyStoreRegistry.getInstance()
                .getDistributedKeyStores()
                .get(storeName);
        return sign(alias, store, ArrayUtils.EMPTY_CHAR_ARRAY);
    }
}
