package kz.gamma;

import kz.gamma.asn1.x509.X509Name;
import kz.gamma.cms.CMSProcessableByteArray;
import kz.gamma.cms.CMSSignedData;
import kz.gamma.cms.CMSSignedDataGenerator;
import kz.gamma.cms.Pkcs7Data;
import kz.gamma.jce.provider.GammaTechProvider;
import kz.gamma.tumarcsp.params.StoreObjectParam;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.security.*;
import java.security.cert.*;
import java.security.cert.Certificate;
import java.util.ArrayList;
import java.util.Date;
import java.util.Enumeration;

/**
 * Created by s_moiseyev
 * Date: 09.10.2009
 * Time: 11:16:26
 */
public class SampleJCECSP_Signature {

    /**
     * @param args
     */
    public static void main(String[] args) {
        Security.addProvider(new GammaTechProvider());
        profileTestPKCS7();
        profileTest();
        profileTest("profile://eToken", "1234567890");
        profileTestPKCS12("C:\\test.p12", "1q2w3e4r5t");
    }

    /**
     *
     */
    public static void profileTest() {
        try {
            //Формируем класс хранилища ключей, будут доступны все профайлы криптопровайдера.
            KeyStore store = loadKeyStore();
            //Получение списка ключей
            Enumeration en = store.aliases();
            while (en.hasMoreElements()) {
                StoreObjectParam profParam = (StoreObjectParam) en.nextElement();
                System.out.println(profParam);
            }
            //Получение закрытого ключа по DN имени сертификата
            PrivateKey prvKey = getPrivateKey("C=KZ,O=kisc,OU=Что за,CN=Моисеев Сергей,SERIALNUMBER=333333333333", store, "");
            if (prvKey != null) {
                byte[] textBlob = "text for sign".getBytes();
                //Постановка подписи
                byte[] sign = SetSign(textBlob, prvKey);
                //Получение сертификата по DN имени
                Certificate cert = getCertificate("C=KZ,O=kisc,OU=Что за,CN=Моисеев Сергей,SERIALNUMBER=333333333333", store, "");
                if (cert != null) {
                    //Проверка подписи
                    if (VerSign(textBlob, sign, cert.getPublicKey()))
                        System.out.println("verify sign OK");
                    else
                        System.out.println("verify sign FALSE");
                }
            }
        } catch (Exception ex) {
            ex.printStackTrace();
        }
    }

    /**
     * @param profile
     * @param pass
     */
    public static void profileTest(String profile, String pass) {
        try {
            //Формируем класс хранилища ключей из профайла.
            KeyStore store = loadKeyStore(profile, pass);
            //Получение списка ключей
            Enumeration en = store.aliases();
            while (en.hasMoreElements()) {
                StoreObjectParam profParam = (StoreObjectParam) en.nextElement();
                System.out.println(profParam);
            }
            //Получение закрытого ключа по DN имени сертификата
            PrivateKey prvKey = getPrivateKey("C=KZ,O=kisc,OU=Что за,CN=Моисеев Сергей,SERIALNUMBER=333333333333", store, pass);
            if (prvKey != null) {
                byte[] textBlob = "text for sign".getBytes();
                //Постановка подписи
                byte[] sign = SetSign(textBlob, prvKey);
                //Получение сертификата по DN имени
                Certificate cert = getCertificate("C=KZ,O=kisc,OU=Что за,CN=Моисеев Сергей,SERIALNUMBER=333333333333", store, pass);
                if (cert != null) {
                    //Проверка подписи
                    if (VerSign(textBlob, sign, cert.getPublicKey()))
                        System.out.println("verify sign OK");
                    else
                        System.out.println("verify sign FALSE");
                }
            }
        } catch (Exception ex) {
            ex.printStackTrace();
        }
    }

    /**
     * @param fileName
     * @param pass
     */
    public static void profileTestPKCS12(String fileName, String pass) {
        try {
            String aliasKey = "";
            //Загрузка хранилища из файла в формате PKCS#12
            KeyStore store = loadKeyStorep12(fileName, pass);
            //Получение списка ключей
            Enumeration en = store.aliases();
            while (en.hasMoreElements()) {
                StoreObjectParam profParam = (StoreObjectParam) en.nextElement();
                System.out.println(profParam);
                aliasKey = profParam.sn;
            }
            //Получение закрытого ключа по SN
            PrivateKey prvKey = (PrivateKey) store.getKey(aliasKey, pass.toCharArray());
            if (prvKey != null) {
                byte[] textBlob = "text for sign".getBytes();
                //Формирование подписи
                byte[] sign = SetSign(textBlob, prvKey);
                //Получение сертификата по SN
                Certificate cert = store.getCertificate(aliasKey);
                if (cert != null) {
                    //Проверка подписи
                    if (VerSign(textBlob, sign, cert.getPublicKey()))
                        System.out.println("verify sign OK");
                    else
                        System.out.println("verify sign FALSE");
                }
            }
        } catch (Exception ex) {
            ex.printStackTrace();
        }
    }

    /**
     *
     */
    public static void profileTestPKCS7() {
        try {
            //Формируем класс хранилища ключей, будут доступны все профайлы криптопровайдера.
            KeyStore store = loadKeyStore();
            //Получение списка ключей
            Enumeration en = store.aliases();
            while (en.hasMoreElements()) {
                StoreObjectParam profParam = (StoreObjectParam) en.nextElement();
                System.out.println(profParam);
            }
            //Получение закрытого ключа по DN имени сертификата
            PrivateKey prvKey = getPrivateKey("C=KZ,O=GAMMA,CN=GAMMAJCE", store, "");
            if (prvKey != null) {
                //Получение сертификата по DN имени
                Certificate cert = getCertificate("C=KZ,O=GAMMA,CN=GAMMAJCE", store, "");
                if (cert != null) {
                    //Сформировать PKCS#7 включая исходный текст
                    byte[] pkcs7 = setSignPKCS7("Message for sign".getBytes(), prvKey, (X509Certificate) cert, true);
                    System.out.println(verSignPKCS7(pkcs7));

                }
            }
        } catch (Exception ex) {
            ex.printStackTrace();
        }
    }

    /**
     * Создаем экземпляр класса для работы с TumarCSP.
     * Данный метод загружает все ключи и сертификаты доступные в данный момент
     *
     * @return
     * @throws NoSuchProviderException
     * @throws KeyStoreException
     * @throws IOException
     * @throws NoSuchAlgorithmException
     * @throws CertificateException
     */
    public static KeyStore loadKeyStore() throws NoSuchProviderException, KeyStoreException, IOException, NoSuchAlgorithmException, CertificateException {
        KeyStore store = KeyStore.getInstance("PKS", "GAMMA");
        store.load(null, null);
        return store;
    }

    /**
     * Создаем экземпляр класса для работы с TumarCSP.
     * Данный метод загружает ключи из выбранного профайла, при этом можно задать пароль на профайл
     *
     * @param profileName
     * @param pass
     * @return
     * @throws NoSuchProviderException
     * @throws KeyStoreException
     * @throws IOException
     * @throws NoSuchAlgorithmException
     * @throws CertificateException
     */
    public static KeyStore loadKeyStore(String profileName, String pass) throws NoSuchProviderException, KeyStoreException, IOException, NoSuchAlgorithmException, CertificateException {
        KeyStore store = KeyStore.getInstance("GKS", "GAMMA");
        store.load(new ByteArrayInputStream(profileName.getBytes()), pass.toCharArray());
        return store;
    }

    /**
     * @param fileName
     * @param pass
     * @return
     * @throws NoSuchProviderException
     * @throws KeyStoreException
     * @throws IOException
     * @throws NoSuchAlgorithmException
     * @throws CertificateException
     */
    public static KeyStore loadKeyStorep12(String fileName, String pass) throws NoSuchProviderException, KeyStoreException, IOException, NoSuchAlgorithmException, CertificateException {
        KeyStore store = KeyStore.getInstance("PKCS12", "GAMMA");
        store.load(new ByteArrayInputStream(fileName.getBytes()), pass.toCharArray());
        return store;
    }

    /**
     * Функция создает экземпляр класса приватного ключа для подписи.
     * Если будет несколько сертификатов с одним именем то загрузить самый новый
     *
     * @param DName
     * @param store
     * @param pass
     * @return
     * @throws KeyStoreException
     * @throws NoSuchAlgorithmException
     * @throws UnrecoverableKeyException
     */
    public static PrivateKey getPrivateKey(String DName, KeyStore store, String pass) throws KeyStoreException, NoSuchAlgorithmException, UnrecoverableKeyException {
        Enumeration en = store.aliases();
        Date tmpDate = null;
        String tmpSN = "";
        while (en.hasMoreElements()) {
            StoreObjectParam prm = (StoreObjectParam) en.nextElement();
            if (prm.dn.length() > 2)
                if ((new X509Name(DName)).equals(new X509Name(prm.dn))) {
                    if (tmpDate == null) {
                        tmpDate = prm.timeCreate;
                        tmpSN = prm.sn;
                    } else {
                        if (prm.timeCreate.after(tmpDate)) {
                            tmpDate = prm.timeCreate;
                            tmpSN = prm.sn;
                        }
                    }
                }
        }
        return (PrivateKey) store.getKey(tmpSN, pass.toCharArray());
    }

    /**
     * Функция создает экземпляр класса сертификата.
     * Если будет несколько сертификатов с одним именем то загрузить самый новый
     *
     * @param DName
     * @param store
     * @param pass
     * @return
     * @throws KeyStoreException
     * @throws NoSuchAlgorithmException
     * @throws UnrecoverableKeyException
     */
    public static Certificate getCertificate(String DName, KeyStore store, String pass) throws KeyStoreException, NoSuchAlgorithmException, UnrecoverableKeyException {
        Enumeration en = store.aliases();
        Date tmpDate = null;
        String tmpSN = "";
        while (en.hasMoreElements()) {
            StoreObjectParam prm = (StoreObjectParam) en.nextElement();
            if (prm.dn.length() > 2)
                if ((new X509Name(DName)).equals(new X509Name(prm.dn))) {
                    if (tmpDate == null) {
                        tmpDate = prm.timeCreate;
                        tmpSN = prm.sn;
                    } else {
                        if (prm.timeCreate.after(tmpDate)) {
                            tmpDate = prm.timeCreate;
                            tmpSN = prm.sn;
                        }
                    }
                }
        }
        return store.getCertificate(tmpSN);
    }

    /**
     * Функция создает цепочку сертификатов.
     * Если будет несколько сертификатов с одним именем то загрузить самый новый
     *
     * @param DName
     * @param store
     * @param pass
     * @return
     * @throws KeyStoreException
     * @throws NoSuchAlgorithmException
     * @throws UnrecoverableKeyException
     */
    public static Certificate[] getCertificateChain(String DName, KeyStore store, String pass) throws KeyStoreException, NoSuchAlgorithmException, UnrecoverableKeyException {
        Enumeration en = store.aliases();
        Date tmpDate = null;
        String tmpSN = "";
        while (en.hasMoreElements()) {
            StoreObjectParam prm = (StoreObjectParam) en.nextElement();
            if (prm.dn.equals(DName)) {
                if (tmpDate == null) {
                    tmpDate = prm.timeCreate;
                    tmpSN = prm.sn;
                } else {
                    if (prm.timeCreate.after(tmpDate)) {
                        tmpDate = prm.timeCreate;
                        tmpSN = prm.sn;
                    }
                }
            }
        }
        return store.getCertificateChain(tmpSN);
    }

    /**
     * Метод, устанавливающий подпись
     *
     * @param buf
     * @param key
     * @return
     */
    public static byte[] SetSign(byte[] buf, PrivateKey key) {
        byte[] sign = null;
        try {
            //Указываем классу Signature что необходимо использовать JCE GAMMA.
            Signature sgn = Signature.getInstance("ECGOST34310", "GAMMA");
            sgn.initSign(key);
            sgn.update(buf);
            sign = sgn.sign();
        }
        catch (Exception ex) {
            ex.printStackTrace();
        }
        return sign;
    }

    /**
     * Метод, проверяющий подпись
     *
     * @param buf
     * @param sign
     * @param key
     * @return
     */
    public static boolean VerSign(byte[] buf, byte[] sign, PublicKey key) {
        boolean ret = false;
        try {
            //Указываем классу Signature что необходимо использовать JCE GAMMA.
            Signature sgn = Signature.getInstance("ECGOST34310", "GAMMA");
            sgn.initVerify(key);
            sgn.update(buf);
            ret = sgn.verify(sign);
        }
        catch (Exception ex) {
            ex.printStackTrace();
        }
        return ret;
    }

    /**
     * Функция ставит подпись в формате pkcs#7
     *
     * @param buf
     * @param key
     * @param cert
     * @param includeBuf
     * @return
     * @throws IOException
     */
    public static byte[] setSignPKCS7(byte[] buf, PrivateKey key, X509Certificate cert, boolean includeBuf) throws IOException {
        CMSSignedData signedData = null;
        try {
            CMSSignedDataGenerator gen = new CMSSignedDataGenerator();
            gen.addSigner(key, cert, CMSSignedDataGenerator.DIGEST_GOST3411G);
            ArrayList list = new ArrayList();
            list.add(cert);
            CertStore store = CertStore.getInstance("Collection", new CollectionCertStoreParameters(list), GammaTechProvider.PROVIDER_NAME);
            gen.addCertificatesAndCRLs(store);
            CMSProcessableByteArray content = new CMSProcessableByteArray(buf);
            signedData = gen.generate(content, includeBuf, GammaTechProvider.PROVIDER_NAME);
        } catch (Exception ex) {
            ex.printStackTrace();
        }
        return signedData.getEncoded();
    }

    /**
     * Функция проверяет подпись в формате pkcs#7
     *
     * @param pkcs7
     * @return
     * @throws IOException
     */
    public static boolean verSignPKCS7(byte[] pkcs7) throws IOException {
        try {
            // Для того чтобы проверить подпись в PKCS#7 без вложеного исходного текста
            // необходимо вызвать Pkcs7Data(pkcs7, буфер с текстом);
            Pkcs7Data signedData = new Pkcs7Data(pkcs7);
            System.out.println(signedData.verify());
        } catch (Exception ex) {
            ex.printStackTrace();
        }
        return true;
    }
}