import kz.gamma.asn1.ASN1Set;
import kz.gamma.asn1.cms.MSTemplate;
import kz.gamma.asn1.cryptopro.KZObjectIndentifiers;
import kz.gamma.asn1.x509.X509Name;
import kz.gamma.cms.CMSSignedData;
import kz.gamma.core.network.RequestSender;
import kz.gamma.jce.PKCS10CertificationRequest;
import kz.gamma.jce.provider.GammaTechProvider;
import kz.gamma.jce.provider.JCEECPrivateKey;
import kz.gamma.jce.provider.JCEECPublicKey;
import kz.gamma.tumarcsp.LibraryWrapper;
import kz.gamma.tumarcsp.TumarCspFunctions;

import java.nio.charset.StandardCharsets;
import java.security.PrivateKey;
import java.security.Security;
import java.security.cert.CertStore;
import java.security.cert.X509CertSelector;
import java.security.cert.X509Certificate;
import java.util.Collection;

public class GenKeyCertTest {
    public enum Alg {
        GOST2004,
        GOST2015,
        RSA2048,
        RSA4096
    }


    public static void main(String... args) {
        try {
            Security.addProvider(new GammaTechProvider());
            //GOST2004:
            genKeyAndCert("profile://test", "C=KZ, O=gamma, CN=user2", "cn=GOST_RAUTIL_USER_1Y,o=Template,c=KZ", "http://172.16.172.233:62310", Alg.GOST2004);

            //GOST2015:
            genKeyAndCert("profile://test", "C=KZ, O=gamma, CN=user2", "CN=GOST_RAUTIL_USER_SIGN_2015_512_OID_1Y;O=Template;C=KZ", "http://172.16.172.65:62310", Alg.GOST2015);

            //RSA2048:
            genKeyAndCert("profile://test", "C=KZ, O=gamma, CN=user2", "CN=Default_CMP;O=Template;C=KZ", "http://172.16.172.65:62310", Alg.RSA2048);

            //RSA4096:
            genKeyAndCert("profile://test", "C=KZ, O=gamma, CN=user2", "CN=Default_CMP;O=Template;C=KZ", "http://172.16.172.65:62310", Alg.RSA4096);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    /**
     * Генерация по указанному профайлу ключа на подпись
     * Формирование PKCS#10 запроса;
     * Отправка запроса в УЦ;
     * Получение ответа от УЦ;
     * Извлечение сертификата из ответа;
     * Импорт сертификата в ключевой контейнер
     *
     * @param profile
     * @param dn
     * @param template
     * @param raUrl
     * @throws Exception
     */
    public static void genKeyAndCert(String profile, String dn, String template, String raUrl, Alg alg) throws Exception {
        String keyOid = null;
        String signatureAlgoritm = null;
        int keyAlg = 0;
        int exportKeyFlag = 0;
        int blobType = 0;
        switch (alg) {
            case GOST2004:
                keyOid = KZObjectIndentifiers.GOST_34_310_KEY_A.getId();
                signatureAlgoritm = LibraryWrapper.KZ_OID_GOST;
                keyAlg = GammaTechProvider.getAlgUtil().algKey.get(keyOid);
                exportKeyFlag = LibraryWrapper.CRYPT_KAZ_EXPORT;
                blobType = LibraryWrapper.PUBLICKEYBLOB_SUBJ_KEY;
                break;
            case GOST2015:
                keyOid = KZObjectIndentifiers.GOSTR_34_10_2015_KEY512_A.getId();
                signatureAlgoritm = LibraryWrapper.KZ_OID_GOST2015_512;
                keyAlg = GammaTechProvider.getAlgUtil().algKey.get(keyOid);
                exportKeyFlag = LibraryWrapper.CRYPT_KAZ_EXPORT;
                blobType = LibraryWrapper.PUBLICKEYBLOB_SUBJ_KEY;
                break;
            case RSA2048:
                keyOid = "1.2.840.113549.1.1.1";
                signatureAlgoritm = "SHA256withRSA";
                keyAlg = 41042;//42066 - sign;
                exportKeyFlag = 0;
                blobType = LibraryWrapper.PUBLICKEYBLOB;
                break;
            case RSA4096:
                keyOid = "1.2.840.113549.1.1.1";
                signatureAlgoritm = "SHA256withRSA";
                keyAlg = 41044;//42068 - sign
                exportKeyFlag = 0;
                blobType = LibraryWrapper.PUBLICKEYBLOB;
                break;
        }

        Number hProv = TumarCspFunctions.cpAcquireContext(profile, LibraryWrapper.CRYPT_NEWKEYSET);
        Number hKey = TumarCspFunctions.cpGenKey(hProv, keyAlg, LibraryWrapper.CRYPT_EXPORTABLE);
        System.out.println("Generate key - success");

        //for PKCS10:
        TumarCspFunctions.cpSetKeyParam(hProv, hKey, LibraryWrapper.KP_USER_KEY, null, 0);

        byte[] publicKey = TumarCspFunctions.cpExportKeyData(hProv, hKey, 0, blobType, exportKeyFlag);
        //System.out.println(UtilCM.array2hex(publicKey));
        System.out.println("Export public key - success");

        JCEECPublicKey pubKey = new JCEECPublicKey(publicKey, keyOid);

        PrivateKey privKey = new JCEECPrivateKey(profile, hProv, hKey);

        //Добавление шаблона
        ASN1Set set = MSTemplate.getASN1Template(template);

        PKCS10CertificationRequest req = new PKCS10CertificationRequest(signatureAlgoritm, new X509Name(dn), pubKey, set, privKey);
        byte[] reqData = RequestSender.sendRequest(raUrl, req.getDEREncoded());
        System.out.println("PKCS#10 send request - success");

        CMSSignedData signedData = new CMSSignedData(reqData);
        CertStore store1 = signedData.getCertificatesAndCRLs("Collection", "GAMMA");
        Collection crtColl = store1.getCertificates(new X509CertSelector());
        if (crtColl.isEmpty() || crtColl.size() == 0) {
            String messageFromCa = "";
            try {
                messageFromCa = new String((byte[]) (signedData.getSignedContent()).getContent(), StandardCharsets.UTF_8);
            } catch (Exception e) {
                //do nothing
            }
            throw new Exception("no user certificate in CA response: " + messageFromCa);
        } else {
            X509Certificate cer = (X509Certificate) crtColl.iterator().next();
            TumarCspFunctions.cpSetKeyParam(hProv, hKey, LibraryWrapper.KP_CERTIFICATE, cer.getEncoded(), LibraryWrapper.ZERO);
            System.out.println("Import certificate to profile - success");
        }
    }
}
