package be.iminds.ilabt.jfed.lowlevel;

import be.iminds.ilabt.jfed.util.KeyUtil;
import be.iminds.ilabt.jfed.util.SSHKeyHelper;
import sun.security.x509.*;

import java.io.IOException;
import java.security.*;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;
import java.util.Date;
import java.util.Objects;
import java.util.UUID;

/**
 * TestCertHelper
 *
 * 2 authorities, a user, a tool and an SA
 *
 * The user, tool and SA are "from" authority 1
 */
public class TestCertHelper {
    public static final int defaultValidityDays = 2;

    public static final String authority1Hostname = "authority.example.com";
    public static final String authority1Urn = "urn:publicid:IDN+" + authority1Hostname + "+authority+root";

    public static final String authority2Hostname = "authority2.example.com";
    public static final String authority2Urn = "urn:publicid:IDN+" + authority2Hostname + "+authority+root";

    public static final String authority3Hostname = "authority3.example.com";
    public static final String authority3Urn = "urn:publicid:IDN+" + authority3Hostname + "+authority+root";

    public static final String fakeHostname = "fake.example.com";

    public static final String username = "tester";
    public static final String userUuid = UUID.randomUUID().toString();
    public static final String userEmail = username + "@" + authority1Hostname;
    public static final String userUuidUrn = "urn:uuid:" + userUuid;
    public static final String userUrn = "urn:publicid:IDN+" + authority1Hostname + "+user+" + username;

    public static final String am1Uuid = UUID.randomUUID().toString();
    public static final String am1Email = "am1@" + authority1Hostname;
    public static final String am1UuidUrn = "urn:uuid:" + am1Uuid;
    public static final String am1Urn = "urn:publicid:IDN+" + authority1Hostname + "+authority+cm";

    public static final String am2Uuid = UUID.randomUUID().toString();
    public static final String am2Email = "am2@" + authority2Hostname;
    public static final String am2UuidUrn = "urn:uuid:" + am2Uuid;
    public static final String am2Urn = "urn:publicid:IDN+" + authority2Hostname + "+authority+cm";

    public static final String am3Uuid = UUID.randomUUID().toString();
    public static final String am3Email = "am3@" + authority3Hostname;
    public static final String am3UuidUrn = "urn:uuid:" + am3Uuid;
    public static final String am3Urn = "urn:publicid:IDN+" + authority3Hostname + "+authority+cm";

    public static final String scsUuid = UUID.randomUUID().toString();
    public static final String scsEmail = "scs@" + authority3Hostname;
    public static final String scsUuidUrn = "urn:uuid:" + scsUuid;
    public static final String scsUrn = am3Urn;//"urn:publicid:IDN+" + authority3Hostname + "+authority+scs";

    public static final String fakeAmUrn = "urn:publicid:IDN+" + fakeHostname + "+authority+cm";

    public static final String saUuid = UUID.randomUUID().toString();
    public static final String saEmail = "sa@" + authority1Hostname;
    public static final String saUuidUrn = "urn:uuid:" + saUuid;
    public static final String saUrn = "urn:publicid:IDN+" + authority1Hostname + "+authority+sa";

    public static final String sa2Uuid = UUID.randomUUID().toString();
    public static final String sa2Email = "sa@" + authority2Hostname;
    public static final String sa2UuidUrn = "urn:uuid:" + sa2Uuid;
    public static final String sa2Urn = "urn:publicid:IDN+" + authority2Hostname + "+authority+sa";

    public static final String toolUuid = UUID.randomUUID().toString();
    public static final String toolUrn = "urn:publicid:IDN+" + authority1Hostname + "+tool+jfed_unit_test";
    public static final String toolUuidUrn = "urn:uuid:" + toolUuid;


    public final SSHKeyHelper authority1Keys;
    public final SSHKeyHelper authority2Keys;
    public final SSHKeyHelper authority3Keys;
    public final SSHKeyHelper userKeys;
    public final SSHKeyHelper saKeys;
    public final SSHKeyHelper am1Keys;
    public final SSHKeyHelper am2Keys;
    public final SSHKeyHelper scsKeys;
    public final SSHKeyHelper toolKeys;

    public final X509Certificate authority1Cert;
    public final X509Certificate authority2Cert;
    public final X509Certificate authority3Cert;
    public final X509Certificate saCert;
    public final X509Certificate am1Cert;
    public final X509Certificate am2Cert;
    public final X509Certificate scsCert;
    public final X509Certificate userCert;
    public final X509Certificate toolCert;
    public final KeyStore testTrustStore;

    public TestCertHelper() {
        try {
            authority1Keys = new SSHKeyHelper();
            authority2Keys = new SSHKeyHelper();
            authority3Keys = new SSHKeyHelper();
            userKeys = new SSHKeyHelper();
            saKeys = new SSHKeyHelper();
            am1Keys = new SSHKeyHelper();
            am2Keys = new SSHKeyHelper();
            scsKeys = new SSHKeyHelper();
            toolKeys = new SSHKeyHelper();

            authority1Cert = makeAuthCert(authority1Keys, authority1Hostname, authority1Urn);
            authority2Cert = makeAuthCert(authority2Keys, authority2Hostname, authority2Urn);
            authority3Cert = makeAuthCert(authority3Keys, authority3Hostname, authority3Urn);
            userCert = makeUserCert(authority1Keys, authority1Hostname, authority1Urn, userKeys, userEmail, userUrn, userUuid, userUuidUrn);
            saCert = makeServerCert(authority1Keys, authority1Hostname, authority1Urn, saKeys, saEmail, saUrn, saUuid, saUuidUrn);
            am1Cert = makeServerCert(authority1Keys, authority1Hostname, authority1Urn, am1Keys, am1Email, am1Urn, am1Uuid, am1UuidUrn);
            am2Cert = makeServerCert(authority1Keys, authority1Hostname, authority1Urn, am2Keys, am2Email, am2Urn, am2Uuid, am2UuidUrn);
            scsCert = makeServerCert(authority3Keys, authority3Hostname, authority3Urn, scsKeys, scsEmail, scsUrn, scsUuid, scsUuidUrn);
            toolCert = makeToolCert(authority1Keys, authority1Hostname, authority1Urn, toolKeys, toolUrn, toolUuid, toolUuidUrn);

            testTrustStore = KeyStore.getInstance("jks");
            testTrustStore.load(null, null);
            testTrustStore.setCertificateEntry("testTrustRoot1", authority1Cert);
            testTrustStore.setCertificateEntry("testTrustRoot2", authority2Cert);
            testTrustStore.setCertificateEntry("testTrustRoot3", authority3Cert);
        } catch (Exception e) {
            throw new RuntimeException("Error generating test keys and/or certificate", e);
        }
    }

//    private void generateCertificates()
//            throws NoSuchAlgorithmException, IOException, CertificateException, SignatureException,
//            NoSuchProviderException, InvalidKeyException, CertPathBuilderException,
//            KeyStoreException, InvalidAlgorithmParameterException {
//
//
////        System.out.println("Authority cert self-signed:\n" + KeyUtil.x509certificateToPem(authorityCert));
////        System.out.println("\nUser cert signed by authority:\n" + KeyUtil.x509certificateToPem(userCert));
////        System.out.println("\nTool cert signed by authority:\n" + KeyUtil.x509certificateToPem(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));
//    }

    public static X509Certificate makeUserCert(SSHKeyHelper authorityKeys, String authHostname, String authUrn,
                                               SSHKeyHelper userKeys, String userEmail, String userUrn, String userUuid, String userUuidUrn) throws NoSuchAlgorithmException, IOException, CertificateException, SignatureException, NoSuchProviderException, InvalidKeyException {
        return makeUserCert(authorityKeys, authHostname, authUrn,
                userKeys, userEmail, userUrn, userUuid, userUuidUrn, defaultValidityDays);
    }
    public static X509Certificate makeUserCert(SSHKeyHelper authorityKeys, String authHostname, String authUrn,
                                               SSHKeyHelper userKeys, String userEmail, String userUrn, String userUuid, String userUuidUrn,
                                               int validityDays) throws NoSuchAlgorithmException, IOException, CertificateException, SignatureException, NoSuchProviderException, InvalidKeyException {
//        Calendar expiry = Calendar.getInstance();
//        expiry.add(Calendar.DAY_OF_YEAR, validityDays);

        PublicKey pubKey = userKeys.getSshPublicKey(); //user public key!
        PrivateKey privKey = authorityKeys.getSshPrivateKey();

        // Determine the signature algorithm
        String sigAlgName = KeyUtil.getCompatibleSigAlgName(privKey.getAlgorithm());

        Date firstDate = new Date(System.currentTimeMillis() - (24L * 60 * 60 * 1000));
        Date lastDate = new Date(System.currentTimeMillis() + (validityDays * 24L * 60 * 60 * 1000));
        CertificateValidity interval = new CertificateValidity(firstDate, lastDate);
        X500Name issuer = new X500Name("CN=" + authHostname);
        X500Name subject = new X500Name("EMAILADDRESS=" + userEmail + ", CN=" + userUuid);


        X509CertInfo info = new X509CertInfo();
        // Add all mandatory attributes
        info.set(X509CertInfo.VERSION,
                new CertificateVersion(CertificateVersion.V3));
        info.set(X509CertInfo.SERIAL_NUMBER, new CertificateSerialNumber(
                new java.util.Random().nextInt() & 0x7fffffff));
        AlgorithmId algID = AlgorithmId.getAlgorithmId(sigAlgName);
        info.set(X509CertInfo.ALGORITHM_ID,
                new CertificateAlgorithmId(algID));
        info.set(X509CertInfo.SUBJECT, subject);
        GeneralNames subjectAltNames = new GeneralNames();
        subjectAltNames.add(new GeneralName(new RFC822Name(userEmail)));
        subjectAltNames.add(new GeneralName(new URIName(userUrn)));
        subjectAltNames.add(new GeneralName(new URIName(userUuidUrn)));
        CertificateExtensions extensions = new CertificateExtensions();
        extensions.set(SubjectAlternativeNameExtension.NAME, new SubjectAlternativeNameExtension(subjectAltNames));
        info.set(X509CertInfo.EXTENSIONS, extensions);
        info.set(X509CertInfo.VALIDITY, interval);
        info.set(X509CertInfo.ISSUER, issuer);

        info.set(X509CertInfo.KEY, new CertificateX509Key(pubKey));

        X509CertImpl cert = new X509CertImpl(info);
        cert.sign(privKey, sigAlgName);

        assert Objects.equals(cert.getSubjectDN().getName(), subject.getName());
        assert Objects.equals(cert.getIssuerDN().getName(), issuer.getName());

        return cert;
    }

    public static X509Certificate makeServerCert(SSHKeyHelper authorityKeys, String authHostname, String authUrn,
                                                 SSHKeyHelper serverKeys, String serverEmail, String serverUrn, String serverUuid, String serverUuidUrn) throws NoSuchAlgorithmException, IOException, CertificateException, SignatureException, NoSuchProviderException, InvalidKeyException {
        return makeServerCert(authorityKeys, authHostname, authUrn,
                serverKeys, serverEmail, serverUrn, serverUuid, serverUuidUrn,
                defaultValidityDays);
    }
    public static X509Certificate makeServerCert(SSHKeyHelper authorityKeys, String authHostname, String authUrn,
                                                 SSHKeyHelper serverKeys, String serverEmail, String serverUrn, String serverUuid, String serverUuidUrn,
                                                 int validityDays) throws NoSuchAlgorithmException, IOException, CertificateException, SignatureException, NoSuchProviderException, InvalidKeyException {
//        Calendar expiry = Calendar.getInstance();
//        expiry.add(Calendar.DAY_OF_YEAR, validityDays);

        PublicKey pubKey = serverKeys.getSshPublicKey(); //sa public key
        PrivateKey privKey = authorityKeys.getSshPrivateKey();

        // Determine the signature algorithm
        String sigAlgName = KeyUtil.getCompatibleSigAlgName(privKey.getAlgorithm());

        Date firstDate = new Date(System.currentTimeMillis() - (24L * 60 * 60 * 1000));
        Date lastDate = new Date(System.currentTimeMillis() + (validityDays * 24L * 60 * 60 * 1000));
        CertificateValidity interval = new CertificateValidity(firstDate, lastDate);
        X500Name issuer = new X500Name("CN=" + authHostname);
        X500Name subject = new X500Name("EMAILADDRESS=" + serverEmail + ", CN=" + serverUuid);


        X509CertInfo info = new X509CertInfo();
        // Add all mandatory attributes
        info.set(X509CertInfo.VERSION,
                new CertificateVersion(CertificateVersion.V3));
        info.set(X509CertInfo.SERIAL_NUMBER, new CertificateSerialNumber(
                new java.util.Random().nextInt() & 0x7fffffff));
        AlgorithmId algID = AlgorithmId.getAlgorithmId(sigAlgName);
        info.set(X509CertInfo.ALGORITHM_ID,
                new CertificateAlgorithmId(algID));
        info.set(X509CertInfo.SUBJECT, subject);
        GeneralNames subjectAltNames = new GeneralNames();
        subjectAltNames.add(new GeneralName(new RFC822Name(saEmail)));
        subjectAltNames.add(new GeneralName(new URIName(saUrn)));
        subjectAltNames.add(new GeneralName(new URIName(saUuidUrn)));
        CertificateExtensions extensions = new CertificateExtensions();
        extensions.set(SubjectAlternativeNameExtension.NAME, new SubjectAlternativeNameExtension(subjectAltNames));
        info.set(X509CertInfo.EXTENSIONS, extensions);
        info.set(X509CertInfo.VALIDITY, interval);
        info.set(X509CertInfo.ISSUER, issuer);

        info.set(X509CertInfo.KEY, new CertificateX509Key(pubKey));

        X509CertImpl cert = new X509CertImpl(info);
        cert.sign(privKey, sigAlgName);

        assert Objects.equals(cert.getSubjectDN().getName(), subject.getName());
        assert Objects.equals(cert.getIssuerDN().getName(), issuer.getName());

        return cert;
    }

    public static X509Certificate makeToolCert(SSHKeyHelper authorityKeys, String authHostname, String authUrn,
                                               SSHKeyHelper toolKeys, String toolUrn, String toolUuid, String toolUuidUrn)
            throws NoSuchAlgorithmException, IOException, CertificateException, SignatureException, NoSuchProviderException, InvalidKeyException {
        return makeToolCert( authorityKeys, authHostname, authUrn,
                toolKeys, toolUrn, toolUuid, toolUuidUrn,
                defaultValidityDays);
    }
    public static X509Certificate makeToolCert(SSHKeyHelper authorityKeys, String authHostname, String authUrn,
                                               SSHKeyHelper toolKeys, String toolUrn, String toolUuid, String toolUuidUrn,
                                               int validityDays)
            throws NoSuchAlgorithmException, IOException, CertificateException, SignatureException, NoSuchProviderException, InvalidKeyException {
//        Calendar expiry = Calendar.getInstance();
//        expiry.add(Calendar.DAY_OF_YEAR, validityDays);

        PublicKey pubKey = toolKeys.getSshPublicKey(); //user public key!
        PrivateKey privKey = authorityKeys.getSshPrivateKey();

        // Determine the signature algorithm
        String sigAlgName = KeyUtil.getCompatibleSigAlgName(privKey.getAlgorithm());

        Date firstDate = new Date(System.currentTimeMillis() - (24L * 60 * 60 * 1000));
        Date lastDate = new Date(System.currentTimeMillis() + (validityDays * 24L * 60 * 60 * 1000));
        CertificateValidity interval = new CertificateValidity(firstDate, lastDate);
        X500Name issuer = new X500Name("CN=" + authHostname);
        X500Name subject = new X500Name("CN=" + toolUuid);


        X509CertInfo info = new X509CertInfo();
        // Add all mandatory attributes
        info.set(X509CertInfo.VERSION,
                new CertificateVersion(CertificateVersion.V3));
        info.set(X509CertInfo.SERIAL_NUMBER, new CertificateSerialNumber(
                new java.util.Random().nextInt() & 0x7fffffff));
        AlgorithmId algID = AlgorithmId.getAlgorithmId(sigAlgName);
        info.set(X509CertInfo.ALGORITHM_ID,
                new CertificateAlgorithmId(algID));
        info.set(X509CertInfo.SUBJECT, subject);
        GeneralNames subjectAltNames = new GeneralNames();
        subjectAltNames.add(new GeneralName(new URIName(toolUrn)));
        subjectAltNames.add(new GeneralName(new URIName(toolUuidUrn)));
        CertificateExtensions extensions = new CertificateExtensions();
        extensions.set(SubjectAlternativeNameExtension.NAME, new SubjectAlternativeNameExtension(subjectAltNames));
        info.set(X509CertInfo.EXTENSIONS, extensions);
        info.set(X509CertInfo.VALIDITY, interval);
        info.set(X509CertInfo.ISSUER, issuer);

        info.set(X509CertInfo.KEY, new CertificateX509Key(pubKey));

        X509CertImpl cert = new X509CertImpl(info);
        cert.sign(privKey, sigAlgName);

        assert Objects.equals(cert.getSubjectDN().getName(), subject.getName());
        assert Objects.equals(cert.getIssuerDN().getName(), issuer.getName());

        return cert;
    }

    /* self signed
    *
    * needs to be CA:
    * X509v3 extensions:
               X509v3 Basic Constraints: critical
                   CA:TRUE
    * */
    public static X509Certificate makeAuthCert(SSHKeyHelper authorityKeys, String authHostname, String authUrn) throws NoSuchAlgorithmException, IOException, CertificateException, SignatureException, NoSuchProviderException, InvalidKeyException {
        return makeAuthCert(authorityKeys, authHostname, authUrn, defaultValidityDays);
    }
    public static X509Certificate makeAuthCert(SSHKeyHelper authorityKeys, String authHostname, String authUrn, int validityDays) throws NoSuchAlgorithmException, IOException, CertificateException, SignatureException, NoSuchProviderException, InvalidKeyException {
//        Calendar expiry = Calendar.getInstance();
//        expiry.add(Calendar.DAY_OF_YEAR, validityDays);

        PublicKey pubKey = authorityKeys.getSshPublicKey();
        PrivateKey privKey = authorityKeys.getSshPrivateKey();

        // Determine the signature algorithm
        String sigAlgName = KeyUtil.getCompatibleSigAlgName(privKey.getAlgorithm());

        Date firstDate = new Date(System.currentTimeMillis() - (24L * 60 * 60 * 1000));
        Date lastDate = new Date(System.currentTimeMillis() + (validityDays * 24L * 60 * 60 * 1000));
        CertificateValidity interval = new CertificateValidity(firstDate, lastDate);
        X500Name issuer = new X500Name("CN=" + authHostname);
        X500Name subject = new X500Name("CN=" + authHostname);


        X509CertInfo info = new X509CertInfo();
        // Add all mandatory attributes
        info.set(X509CertInfo.VERSION,
                new CertificateVersion(CertificateVersion.V3));
        info.set(X509CertInfo.SERIAL_NUMBER, new CertificateSerialNumber(
                new java.util.Random().nextInt() & 0x7fffffff));
        AlgorithmId algID = AlgorithmId.getAlgorithmId(sigAlgName);
        info.set(X509CertInfo.ALGORITHM_ID, new CertificateAlgorithmId(algID));

        info.set(X509CertInfo.SUBJECT, subject);
        GeneralNames subjectAltNames = new GeneralNames();
        subjectAltNames.add(new GeneralName(new URIName(authUrn)));

           /* depth of the certification path
               = max number of "hops" to end entity cert.
               1 means this can issue a CA certificate, but that CA cannot issue a CA certificate itself
            */
        BasicConstraintsExtension basicConstraintsExtension =
                new BasicConstraintsExtension(true/*is CA*/, 1 /*depth of the certification path*/);

        CertificateExtensions extensions = new CertificateExtensions();
        extensions.set(SubjectAlternativeNameExtension.NAME, new SubjectAlternativeNameExtension(subjectAltNames));
        extensions.set(BasicConstraintsExtension.NAME, basicConstraintsExtension);

        info.set(X509CertInfo.EXTENSIONS, extensions);
        info.set(X509CertInfo.VALIDITY, interval);
        info.set(X509CertInfo.ISSUER, issuer);

        info.set(X509CertInfo.KEY, new CertificateX509Key(pubKey));

        X509CertImpl cert = new X509CertImpl(info);
        cert.sign(privKey, sigAlgName);

        return cert;
    }

    public String getUserPem() {
        String privKeyPem = new String(userKeys.getPEMAnyPrivateKey());
        String certPem = KeyUtil.x509certificateToPem(userCert);
        return privKeyPem+"\n"+certPem;
    }
}
