package be.iminds.ilabt.jfed.lowlevel;

import be.iminds.ilabt.jfed.util.KeyUtil;
import org.apache.xml.security.Init;
import org.apache.xml.security.transforms.Transforms;
import org.apache.xml.security.utils.Constants;
import org.apache.xml.security.utils.ElementProxy;
import org.testng.Assert;
import org.testng.annotations.Test;
import org.w3c.dom.Document;
import org.w3c.dom.Element;

import javax.xml.crypto.dsig.*;
import javax.xml.crypto.dsig.dom.DOMSignContext;
import javax.xml.crypto.dsig.keyinfo.KeyInfo;
import javax.xml.crypto.dsig.keyinfo.KeyInfoFactory;
import javax.xml.crypto.dsig.keyinfo.KeyValue;
import javax.xml.crypto.dsig.keyinfo.X509Data;
import javax.xml.crypto.dsig.spec.C14NMethodParameterSpec;
import javax.xml.crypto.dsig.spec.TransformParameterSpec;
import javax.xml.transform.OutputKeys;
import javax.xml.transform.Transformer;
import javax.xml.transform.TransformerFactory;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.stream.StreamResult;
import java.io.IOException;
import java.io.StringWriter;
import java.security.*;
import java.security.cert.CertPathBuilderException;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.util.*;

/**
 * CredentialTests:
 * <p/>
 * various tests of credentials.
 * This creates fake credentials, for an authority, a user and a tool.
 * It creates speaksFor and Delegated credentials for both.
 * It tests the validity of all credentials.
 *
 * @link http://groups.geni.net/geni/wiki/GeniApiCertificates
 * @link http://groups.geni.net/geni/wiki/GAPI_AM_API_DRAFT/SpeaksFor
 * @see be.iminds.ilabt.jfed.util.KeyUtil
 * @see be.iminds.ilabt.jfed.util.SSHKeyHelper
 * <p/>
 * Geni requirements for Certificates
 */
public class CredentialTests {
    private static org.slf4j.Logger LOG = org.slf4j.LoggerFactory.getLogger(CredentialTests.class);

    private SfaCredential userCredential; //signed by authority
    private SfaCredential toolCredential; //signed by tool itself or other authority (to be decided)

    private TestCertHelper testCertHelper;

    @Test
    public void generateCertificates()
            throws NoSuchAlgorithmException, IOException, CertificateException, SignatureException,
            NoSuchProviderException, InvalidKeyException, CertPathBuilderException,
            KeyStoreException, InvalidAlgorithmParameterException {
        testCertHelper = new TestCertHelper();

        System.out.println("Authority cert self-signed:\n" + KeyUtil.x509certificateToPem(testCertHelper.authority1Cert));
        System.out.println("\nUser cert signed by authority:\n" + KeyUtil.x509certificateToPem(testCertHelper.userCert));
        System.out.println("\nTool cert signed by authority:\n" + KeyUtil.x509certificateToPem(testCertHelper.toolCert));

        //can't get this check to work in java 8. But manual check using openssl (openssl verify -CAfile ...) conformed the certs are still ok
        //        Assert.assertTrue(checkCert(authorityCert, userCert));
        //        Assert.assertTrue(checkCert(authorityCert, toolCert));
    }

     //can't get this check to work in java 8.
//    public boolean checkCert(X509Certificate trustedAuthority, X509Certificate certToCheck)
//            throws KeyStoreException, InvalidAlgorithmParameterException, NoSuchAlgorithmException, CertPathBuilderException {
//        List<X509Certificate> chain = new ArrayList<X509Certificate>();
//        chain.add(certToCheck);
//
//        /* Construct a valid path. */
//        KeyStore anchors = testTrustStore;//JFedTrustStoreHelper.getFullTrustStore();
//        anchors.setCertificateEntry("test-authority", trustedAuthority);
//
//        X509CertSelector target = new X509CertSelector();
//        target.setCertificate(chain.get(0));
//        PKIXBuilderParameters params = new PKIXBuilderParameters(anchors, target);
//        CertStoreParameters intermediates = new CollectionCertStoreParameters(chain);
//        params.addCertStore(CertStore.getInstance("Collection", intermediates));
////        CertStoreParameters revoked = new CollectionCertStoreParameters(crls);
////        params.addCertStore(CertStore.getInstance("Collection", revoked));
//        CertPathBuilder builder = CertPathBuilder.getInstance("PKIX");
//        /*
//         * If build() returns successfully, the certificate is valid. More details
//         * about the valid path can be obtained through the PKIXBuilderResult.
//         * If no valid path can be found, a CertPathBuilderException is thrown.
//         */
//        CertPathBuilderResult r = (CertPathBuilderResult) builder.build(params);
//        return r != null && r.getCertPath() != null;
//    }


    public static String signXml(Document doc, String elementToSignId, String elementToAddSignatureToName, X509Certificate cert, Key privateKey) throws Exception {
        doc.normalizeDocument();

        if (cert.getPublicKey() instanceof RSAPublicKey && privateKey instanceof RSAPrivateKey) {
            assert Objects.equals(((RSAPublicKey) cert.getPublicKey()).getModulus(), ((RSAPrivateKey) privateKey).getModulus());
            if (LOG.isDebugEnabled()) {
                boolean same =
                        Objects.equals(((RSAPublicKey) cert.getPublicKey()).getModulus(), ((RSAPrivateKey) privateKey).getModulus());
                LOG.debug("signXml input cert and privateKey have same modulus");
            }
        } else
            LOG.debug("not Rsa private and/or public key: no check done");


//        return signXmlWithoutSanturio(doc, elementToSignId, elementToAddSignatureToName, cert, privateKey);
        return signXmlWithSanturio(doc, elementToSignId, elementToAddSignatureToName, cert, privateKey);
    }

    public static String signXmlWithoutSanturio(Document doc, String elementToSignId, String elementToAddSignatureToName, X509Certificate cert, Key privateKey)
            throws Exception {
        Element elementToAddSignatureTo = (Element) doc.getElementsByTagName(elementToAddSignatureToName).item(0);
//        Element elementToSign = doc.getElementById(elementToSignId);
        assert elementToAddSignatureTo != null;
//        assert elementToSign != null;

        XMLSignatureFactory fac = XMLSignatureFactory.getInstance("DOM");

        Reference ref = fac.newReference
                ("#" + elementToSignId, fac.newDigestMethod(DigestMethod.SHA1, null),
                        Collections.singletonList
                                (fac.newTransform
                                        (Transform.ENVELOPED, (TransformParameterSpec) null)),
                        null, null
                );

        // Create the SignedInfo
        SignedInfo si = fac.newSignedInfo
                (fac.newCanonicalizationMethod
                                (CanonicalizationMethod.INCLUSIVE,
                                        (C14NMethodParameterSpec) null),
                        fac.newSignatureMethod(SignatureMethod.RSA_SHA1, null),
                        Collections.singletonList(ref)
                );


        // Create a KeyValue
        KeyInfoFactory kif = fac.getKeyInfoFactory();
        KeyValue kv = kif.newKeyValue(cert.getPublicKey());
        List<X509Certificate> x509DataList = new ArrayList<>();
        x509DataList.add(cert); //adds X509Certificate element
//        x509DataList.add(""); //adds X509SubjectName element
        //can also use X509IssuerSerial, see newX509Data javadoc
        X509Data x509Data = kif.newX509Data(x509DataList);

        // Create a KeyInfo and add the KeyValue to it
        List<Object> kiList = new ArrayList<>();
        kiList.add(kv);
        kiList.add(x509Data);
        KeyInfo ki = kif.newKeyInfo(/*Collections.singletonList(kv)*/kiList);


        // Create a DOMSignContext and specify the PrivateKey and
        // location of the resulting XMLSignature's parent element
        DOMSignContext dsc = new DOMSignContext(privateKey, elementToAddSignatureTo);

        // Create the XMLSignature (but don't sign it yet)
        javax.xml.crypto.dsig.XMLSignature signature = fac.newXMLSignature(si, ki);

        // Marshal, generate (and sign) the enveloped signature
        signature.sign(dsc);


        TransformerFactory tf = TransformerFactory.newInstance();
        Transformer trans = tf.newTransformer();
        trans.setOutputProperty(OutputKeys.OMIT_XML_DECLARATION, "no");
        trans.setOutputProperty(OutputKeys.ENCODING, "utf-8");
        trans.setOutputProperty(OutputKeys.INDENT, "no");
        final StringWriter stringWriter = new StringWriter();
        trans.transform(new DOMSource(doc), new StreamResult(stringWriter));
        return stringWriter.getBuffer().toString();

    }

    public static String signXmlWithSanturio(Document doc, String elementToSignId, String elementToAddSignatureToName, X509Certificate cert, Key privateKey)
            throws Exception {
        Init.init();
        ElementProxy.setDefaultPrefix(Constants.SignatureSpecNS, "");
        final org.apache.xml.security.signature.XMLSignature sig =
                new org.apache.xml.security.signature.XMLSignature(doc, null, org.apache.xml.security.signature.XMLSignature.ALGO_ID_SIGNATURE_RSA);
//        final XMLSignature sig = new XMLSignature(elementToSign, null);
        Element elementToAddSignatureTo = (Element) doc.getElementsByTagName(elementToAddSignatureToName).item(0);
        elementToAddSignatureTo.appendChild(sig.getElement());
        final Transforms transforms = new Transforms(doc);
        transforms.addTransform(Transforms.TRANSFORM_ENVELOPED_SIGNATURE);
//        transforms.addTransform(Transforms.TRANSFORM_C14N_OMIT_COMMENTS);
        sig.addDocument("#" + elementToSignId, transforms, Constants.ALGO_ID_DIGEST_SHA1);
        sig.addKeyInfo(cert);
        sig.addKeyInfo(cert.getPublicKey());
        sig.sign(privateKey);


        final StringWriter stringWriter = new StringWriter();
        TransformerFactory transFactory = TransformerFactory.newInstance();
        Transformer transformer = transFactory.newTransformer();
        transformer.setOutputProperty(OutputKeys.OMIT_XML_DECLARATION, "no");
        transformer.setOutputProperty(OutputKeys.ENCODING, "utf-8");
        transformer.setOutputProperty(OutputKeys.INDENT, "no");
        transformer.transform(new DOMSource(doc), new StreamResult(stringWriter));
        return stringWriter.getBuffer().toString();


        //meses up xsi:xmlns
//        ByteArrayOutputStream os = new ByteArrayOutputStream();
//        os.write(Canonicalizer.getInstance(Canonicalizer.ALGO_ID_C14N_WITH_COMMENTS).canonicalizeSubtree(doc));
//        return os.toString("UTF-8");
    }


//    @Test
//    public void checkRealCredentials() {
//        logger.debug("Checking w Credential");
//        SfaCredential wCred = new SfaCredential("w-sa-cred", wCredential, "geni_sfa", "3");
//
//        Assert.assertTrue(wCred.check());
//
//        logger.debug("Checking f Credential");
//        SfaCredential fCred = new SfaCredential("f-sa-cred", fCredential, "geni_sfa", "3");
//
//        Assert.assertTrue(fCred.check());
//    }

    @Test(dependsOnMethods = {"generateCertificates" })
    public void createUserCredential()
            throws Exception {
        Date expireDate = new Date(System.currentTimeMillis() + (2 * 24 * 60 * 60 * 1000)); //2 days from now

        LOG.debug("Creating User Credential");

        userCredential =
                SfaCredential.create("privilege", TestCertHelper.userUrn,
                        TestCertHelper.saUrn, testCertHelper.userCert, testCertHelper.saCert, testCertHelper.saCert,
                        testCertHelper.saKeys.getSshPrivateKey(), expireDate, "*", false, "Test User Credential");
        System.out.println("User credential:\n" + userCredential.getCredentialXml() + "\n");
        LOG.trace("User credential:\n" + userCredential.getCredentialXml());
    }

    @Test(dependsOnMethods = {"createUserCredential"})
    public void checkUserCredential() throws KeyStoreException, CredentialException {
        LOG.debug("Checking User Credential");

        KeyStore trustStore = testCertHelper.testTrustStore;//JFedTrustStoreHelper.getFullTrustStore();
        trustStore.setCertificateEntry("test-authority", testCertHelper.authority1Cert);
        Assert.assertTrue(userCredential.check(trustStore));

        LOG.debug("User Credential check passed");
    }

    @Test(dependsOnMethods = {"checkUserCredential"})
    public void createAndCheckToolCredential()
            throws Exception {
        Date expireDate = new Date(System.currentTimeMillis() + (2 * 24 * 60 * 60 * 1000)); //2 days from now
        toolCredential =
                SfaCredential.create("privilege", TestCertHelper.authority1Urn, TestCertHelper.toolUrn, testCertHelper.authority1Cert, testCertHelper.toolCert, testCertHelper.authority1Cert,
                        testCertHelper.authority1Keys.getSshPrivateKey(), expireDate, "*", false, "Test Tool Credential");

        KeyStore trustStore = testCertHelper.testTrustStore;//JFedTrustStoreHelper.getFullTrustStore();
        trustStore.setCertificateEntry("test-authority", testCertHelper.authority1Cert);
        Assert.assertTrue(toolCredential.check(trustStore));
    }

    SfaCredential speaksForCred;

    @Test(dependsOnMethods = {"createAndCheckToolCredential"})
    public void createSpeaksForCredential() throws Exception {
        Date expireDate = new Date(System.currentTimeMillis() + (3 * 12 * 60 * 60 * 1000)); //1.5 days from now
        speaksForCred =
                SfaCredential.createSpeaksFor(TestCertHelper.userUrn, TestCertHelper.toolUrn, testCertHelper.userCert,
                        testCertHelper.toolCert, testCertHelper.userKeys.getSshPrivateKey(), expireDate, "*", false);
    }

    @Test(dependsOnMethods = {"createSpeaksForCredential"})
    public void checkSpeaksForCredential() throws Exception {
        assert speaksForCred != null;
        LOG.debug("Checking Speaks For credential");
        Assert.assertTrue(speaksForCred.check(testCertHelper.testTrustStore));
    }

    SfaCredential delegatedCred;

    @Test(dependsOnMethods = {"checkSpeaksForCredential"})
    public void createDelegatedCredential() throws Exception {


        Date expireDate = new Date(System.currentTimeMillis() + (3 * 12 * 60 * 60 * 1000)); //1.5 days from now
        delegatedCred =
                userCredential.delegate(TestCertHelper.toolUrn, testCertHelper.toolCert,
                        testCertHelper.userKeys.getSshPrivateKey(), expireDate, "info", false);
    }

    @Test(dependsOnMethods = {"createDelegatedCredential"})
    public void checkDelegatedCredential() throws CredentialException {
        assert delegatedCred != null;
        LOG.debug("Checking Delegated credential");
        Assert.assertTrue(delegatedCred.check(testCertHelper.testTrustStore));
    }


    //below are only predefined string needed for the tests


//    private static String fCredential = ...
//    private static String wCredential= ...
}
