import kz.gamma.asn1.DERGeneralizedTime;
import kz.gamma.asn1.DERObjectIdentifier;
import kz.gamma.asn1.DERPrintableString;
import kz.gamma.asn1.DERSet;
import kz.gamma.asn1.cms.Attribute;
import kz.gamma.asn1.cms.EnvelopedData;
import kz.gamma.cms.CMSGammaEnvelopedData;
import kz.gamma.core.UtilCM;
import kz.gamma.jce.provider.GammaTechProvider;
import kz.gamma.jce.provider.JCEECPrivateKey;
import kz.gamma.tumarcsp.LibraryWrapper;
import kz.gamma.tumarcsp.TumarCspFunctions;
import kz.gamma.util.encoders.Base64;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.security.Security;
import java.security.cert.CertificateException;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.util.Arrays;
import java.util.Date;

public class EncryptHsmTest {
    public static void main(String[] args) throws Exception {
        Security.addProvider(new GammaTechProvider());
        long start;
        byte[] secretText = "secret text".getBytes();

        boolean ENCRYPT = true;
        if (args.length != 0) {
            //в случае получения любого параметра - выполняем decrypt
            ENCRYPT = false;
        }

        if (ENCRYPT) {
            System.out.println("Encrypt start ...");
            String profileSender = "afm_hsm_user1";
            X509Certificate recipientCert = UtilCM.loadCertFromFile("/test/cert_fs_user.cer");
            String secretDataBase64 = new String(Base64.encode(secretText));

            //шифрование:
            start = System.currentTimeMillis();
            String encryptedBase64 = encryptData(secretDataBase64, profileSender, new X509Certificate[]{recipientCert});
            System.out.println(String.format("encrypted (time: %d ms): " + encryptedBase64, System.currentTimeMillis() - start));
        } else {
            //расшифрование:
            System.out.println("Decrypt start ...");
            String encryptedBase64 = "MIAGCSqGSIb3DQEHA6CCBUQwggVAAgECoIIC6qCCAuYwggLiMIICSqADAgECAiAzxAUHfwQr+78bwFdBe5Lli2Ti7VMDM0PmFKw5wuNe6DAOBgorBgEEAbURAQIHBQAwNDELMAkGA1UEBhMCS1oxEDAOBgNVBAoTBzQgTEVWRUwxEzARBgNVBAMTCjQgTEVWRUwgQ0EwHhcNMjUxMTE5MTAxMDM0WhcNMjYxMTE5MTAxMDM0WjA2MQswCQYDVQQGEwJLWjEQMA4GA1UEChMHQUZNVEVTVDEVMBMGA1UEAxMMVVNFUl9IU01fQUZNMIGkMA4GCisGAQQBtREBCAsFAAOBkQAGAgAATqAAAABFQzEABAAAkfJ8+pLg1OCwcnZUjMITq/T+aCFveHomT64SaQ0/1Lp6FUi4MEpy4OrkpMhwHKc\n" +
                    "BpiWQm9RZROuXhBxJ3XWe6FZV3Lwesoad8U3JtOYCtuU/MjYdEV9/uK/wv/u8Z49Tb1q4ArOW3/7TzjTowdnuN3bINmV12K9atnhKkq0+/8OjgdswgdgwCwYDVR0PBAQDAgQwMBMGA1UdJQQMMAoGCCsGAQUFBwMEMCkGA1UdDgQiBCAzxAUHfwQr+78bwFdBe5Lli2Ti7VMDM0PmFKw5wuNe6DCBiAYDVR0jBIG\n" +
                    "AMH6AIEwyDQV/zxpjGKNMLyxoKV4YHbqOzVLHdWuXIf/HxY5toTikNjA0MQswCQYDVQQGEwJLWjEQMA4GA1UEChMHNCBMRVZFTDETMBEGA1UEAxMKNCBMRVZFTCBDQYIgTDINBX/PGmMYo0wvLGgpXhgduo7NUsd1a5ch/8fFjm0wDgYKKwYBBAG1EQECBwUAA4GBAFTlEfFYkFWdtomHPMFXKg1RTwF6lMQ13ig\n" +
                    "9ik6k0DG8bOhq4C8q10+qSJRwIfgHUzG/uDNQFAdjDw9U7lK/ft1tYiHdD8ixzXc4UXeudEsx1PxvAgJSLaBxtSbuL+sivc+gh/vNlrGvB9gn2fjXt8wKZSXbqiu58mf/AC6Tevo5MYIBPDCCATgCAQOgWjBYMDQxEzARBgNVBAMMCjQgTEVWRUwgQ0ExEDAOBgNVBAoMBzQgTEVWRUwxCzAJBgNVBAYTAktaAiA\n" +
                    "zxAUHfwQr+78bwFdBe5Lli2Ti7VMDM0PmFKw5wuNe6DAOBgorBgEEAbURAQIEBQAwgcYwgcMwWDA0MQswCQYDVQQGEwJLWjEQMA4GA1UECgwHNCBMRVZFTDETMBEGA1UEAwwKNCBMRVZFTCBDQQIgGYSVf9DMpqMzp/A8cJ46Nd9sXJwe6CoKh0ZWnKR/mVEEZwECAAAfaAQAH2gAADBZAgEEBAjB4HXSaeWNRgQ\n" +
                    "IydOjGDk75sEEQA94MV2Qh1iu3I8G3pGxKRRMtwfrAkyeH9fzjNfGlLxBUlB/0LrGW/i06GnikcFlYWsZV+XlzDXD/UUigkfFUucwKAYJKoZIhvcNAQcCMA4GCisGAQQBtREBAgQFAIALMo/BTlxsQM4RDNChgeIwGAYJKoZIhvcNAQkDMQsGCSqGSIb3DQEHAzAZBgkqhkiG9w0BCQ0xDBMKREVTQ1JJUFRPUjA\n" +
                    "eBgkqhkiG9w0BCQUxERgPMjAyNTExMTkxMDM4MzZaMEQGCSsGAQQBtRECCTE3MDUxCzAJBgNVBAYTAktaMRAwDgYDVQQKDAdBRk1URVNUMRQwEgYDVQQDDAtVU0VSX0ZTX0FGTTBFBgkrBgEEAbURAggxODA2MRUwEwYDVQQDDAxVU0VSX0hTTV9BRk0xEDAOBgNVBAoMB0FGTVRFU1QxCzAJBgNVBAYTAktaAAA=";
            String profileRecipient = "user_GOST2012ru";

            start = System.currentTimeMillis();
            String decryptedBase64 = decryptText(encryptedBase64, profileRecipient);
            byte[] decryptedText = Base64.decode(decryptedBase64);
            System.out.println(String.format("decrypted: (time: %d ms): " + new String(decryptedText), System.currentTimeMillis() - start));

            if (!Arrays.equals(secretText, decryptedText)) {
                throw new Exception("Decryption failed");
            }
        }
    }

    /**
     * @param dataBase64         - данные в base64
     * @param profileSender      - профайл отправителя
     * @param profilesRecipients - сертификаты получателей
     * @return вернется зашифрованный пакет в base64
     * @throws IOException
     * @throws CertificateException
     */
    public static String encryptData(String dataBase64, String profileSender, X509Certificate[] profilesRecipients) throws IOException, CertificateException {
        try (JCEECPrivateKey prvKey = getKey(profileSender)) {
            X509Certificate certSender = getCert(profileSender);
            byte[] text = Base64.decode(dataBase64);

            CMSGammaEnvelopedData data = new CMSGammaEnvelopedData();
            data.cryptText(text, prvKey);

            for (int i = 0; i < profilesRecipients.length; i++) {
                data.addKeyRecipient(profilesRecipients[i]);
            }

            data.addUnprotectedAttrs(new Attribute((new DERObjectIdentifier("1.2.840.113549.1.9.3")), new DERSet(new DERObjectIdentifier("1.2.840.113549.1.7.3"))));
            data.addUnprotectedAttrs(new Attribute((new DERObjectIdentifier("1.2.840.113549.1.9.5")), new DERSet(new DERGeneralizedTime(new Date()))));
            data.addUnprotectedAttrs(new Attribute((new DERObjectIdentifier("1.2.840.113549.1.9.13")), new DERSet(new DERPrintableString("DESCRIPTOR"))));
            EnvelopedData env = data.generateEnvelopedData(certSender);
            return new String(Base64.encode(env.getEncoded()));
        }
    }

    /**
     * @param encryptedBase64
     * @param profileRecipient - профайл получателя
     * @return вернется расшифрованный пакет в base64
     * @throws Exception
     */
    public static String decryptText(String encryptedBase64, String profileRecipient) throws Exception {
        byte[] encrypted = Base64.decode(encryptedBase64);
        CMSGammaEnvelopedData cmsEn = new CMSGammaEnvelopedData(encrypted);
        X509Certificate certificate = getCert(profileRecipient);
        try (JCEECPrivateKey key = getKey(profileRecipient)) {
            return new String(Base64.encode(cmsEn.decryptText(key, certificate.getSerialNumber())));
        }
    }

    private static JCEECPrivateKey getKey(String profile) {
        Number hProv = TumarCspFunctions.cpAcquireContext(getFullProfileName(profile), 0);
        JCEECPrivateKey privateKey = new JCEECPrivateKey(profile, hProv, null);
        return privateKey;
    }

    private static X509Certificate getCert(String profile) throws CertificateException {
        Number hProv = TumarCspFunctions.cpAcquireContext(getFullProfileName(profile), 0);
        Number hKey = TumarCspFunctions.cpGetUserKey(hProv, LibraryWrapper.AT_KEYEXCHANGE);
        byte[] cert = TumarCspFunctions.cpGetKeyParamData(hProv, hKey, LibraryWrapper.KP_CERTIFICATE, LibraryWrapper.ZERO);
        CertificateFactory cf = CertificateFactory.getInstance("X.509");
        X509Certificate x509cert = (X509Certificate) cf.generateCertificate(new ByteArrayInputStream(cert));
        TumarCspFunctions.cpReleaseContext(hProv, 0);
        return x509cert;
    }

    private static String getFullProfileName(String containerName) {
        return "profile://" + containerName;
    }
}
