import kz.gamma.asn1.cryptopro.KZObjectIndentifiers;
import kz.gamma.asn1.pkcs.PKCSObjectIdentifiers;
import kz.gamma.jce.provider.GammaTechProvider;
import kz.gamma.jce.provider.X509CertificateObject;
import kz.gamma.tumarcsp.LibraryWrapper;
import kz.gamma.tumarcsp.TumarCspFunctions;
import kz.gamma.tumarcsp.params.StoreObjectParam;
import kz.gamma.xmldsig.JCPXMLDSigInit;
import org.apache.xml.security.algorithms.MessageDigestAlgorithm;
import org.apache.xml.security.encryption.XMLCipherParameters;
import org.apache.xml.security.keys.KeyInfo;
import org.apache.xml.security.signature.XMLSignature;
import org.apache.xml.security.transforms.Transforms;
import org.apache.xml.security.utils.Constants;
import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.w3c.dom.NodeList;
import org.xml.sax.InputSource;

import javax.xml.parsers.DocumentBuilder;
import javax.xml.parsers.DocumentBuilderFactory;
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.*;
import java.security.*;
import java.security.cert.Certificate;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;
import java.util.Enumeration;

public class SampleJCE_SignatureXML {


    public static void main(String[] args) {
        try {
            Security.addProvider(new GammaTechProvider());

            System.setProperty("org.apache.xml.security.ignoreLineBreaks", "true");
            org.apache.xml.security.Init.init();

            JCPXMLDSigInit.init();
            String pass = "";
            //Gost 2004
            String profile = createProfile("*","D:\\*",pass,"bin");
            String sourceXML = "D:\\*\\test.xml";
            String targetXML = "D:\\*\\test_sign.xml";
            //            testSign2004(profile, pass,sourceXML, targetXML);
            testSign(profile, pass,sourceXML, targetXML);

            pass = "";
            //Gost 2015
            profile = createProfile("*","D:\\*\\test_2015",pass,"bin");
            targetXML = "D:\\*\\test_sign2.xml";
            //здесь выбираем алгоритм подписи исходя из сертификата
            testSign(profile, pass, sourceXML, targetXML);


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

    }

    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;
    }
    static StoreObjectParam getSignAlias(KeyStore store) throws KeyStoreException {
        Enumeration en = store.aliases();
        StoreObjectParam alias = null;
        while (en.hasMoreElements()) {
            StoreObjectParam profParam = (StoreObjectParam) en.nextElement();
            System.out.println(profParam);
            try{
                X509Certificate cert = profParam.getX509certificate();
                //находим ключ для подписи
                if(alias == null && cert != null && cert instanceof X509CertificateObject && cert.getKeyUsage()[0]  ) {
                    if(cert.getNotAfter().after(new java.util.Date())) {
                        alias = profParam;
                        break;
                    }
                }

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

    public static Document sign2004XML(Document doc, Certificate cert, PrivateKey privKey)
            throws Exception {
        String signMethod = "http://www.w3.org/2001/04/xmldsig-more#gost34310-gost34311";
        String digestMethod = "http://www.w3.org/2001/04/xmldsig-more#gost34311";
//        String signMethod = Constants.MoreAlgorithmsSpecNS+"gost34310-gost34311";
//        String digestMethod = Constants.MoreAlgorithmsSpecNS+"gost34311";

        XMLSignature sig = new XMLSignature(doc, "", signMethod);

        String res = "";
        if (doc.getFirstChild() != null) {
            doc.getFirstChild().appendChild(sig.getElement());
            Transforms transforms = new Transforms(doc);
            transforms.addTransform("http://www.w3.org/2000/09/xmldsig#enveloped-signature");
            transforms.addTransform("http://www.w3.org/TR/2001/REC-xml-c14n-20010315#WithComments");
            sig.addDocument("", transforms, digestMethod);
            sig.addKeyInfo((X509Certificate) cert);
            sig.sign(privKey);
            StringWriter os = new StringWriter();
            TransformerFactory tf = TransformerFactory.newInstance();
            Transformer trans = tf.newTransformer();
            trans.transform(new DOMSource(doc), new StreamResult(os));
            os.flush();
            res = os.toString();
            os.close();
        }
        return parseXml(res);
    }

    public static Document parseXml(String xml)
            throws Exception {
        DocumentBuilderFactory dbf = DocumentBuilderFactory.newInstance();
        dbf.setNamespaceAware(true);
        DocumentBuilder builder = dbf.newDocumentBuilder();
        InputSource source = null;
        source = new InputSource(new StringReader(xml));
        return builder.parse(source);
    }
    public static void testSign2004(String profile, String pass, String filename, String signFilename) {
        try {
            // Формируем класс хранилища ключей, будут доступны все профайлы криптопровайдера.
            KeyStore store = loadKeyStore(profile,pass);
            //Получение списка ключей
            Enumeration en = store.aliases();
            StoreObjectParam alias = getSignAlias(store);
            if(alias == null)
                throw new Exception("Sign key not found");
            // Получение закрытого ключа по DN имени сертификата
//            PrivateKey prvKey = getPrivateKey("C=KZ, O=AutoTest, CN=Test1521", store, "");
            PrivateKey prvKey = (PrivateKey)store.getKey(alias.getSn(), pass.toCharArray()) ;
            if (prvKey != null) {
                //Получение сертификата по DN имени
//                Certificate cert = getCertificate("C=KZ, O=AutoTest, CN=Test1521", store, "");
                Certificate cert = alias.getX509certificate();
                if (cert != null) {
                    DocumentBuilderFactory dbf = DocumentBuilderFactory.newInstance();
                    dbf.setNamespaceAware(true);
                    DocumentBuilder builder = dbf.newDocumentBuilder();
                    InputSource source = null;
                    source = new InputSource(new FileInputStream(filename));
                    // Подписываем XML документ
                    Document sigDoc = sign2004XML(builder.parse(source), cert, prvKey);
                    // Проверяем подпись XML документа
                    if (!validateXML(sigDoc))
                        throw new Exception("Подпись не прошла проверку");
                    else{
                        StringWriter os = new StringWriter();
                        Transformer trans  = TransformerFactory.newInstance().newTransformer();
                        trans.transform(new DOMSource(sigDoc), new StreamResult(os));
                        os.flush();
                        String res = os.toString().replaceAll("&#13;","");
                        os.close();
                        FileWriter fs = new FileWriter(signFilename);
                        fs.write(res);
                        fs.close();
                    }

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

    /**
     * Подпись с выбором Digest исходя из алгоритма подписи сертификата
     * @param doc
     * @param cert
     * @param privKey
     * @return
     * @throws Exception
     */
    public static Document signXML(Document doc, X509Certificate cert, PrivateKey privKey)
            throws Exception {
        String sigAlgOid = cert.getSigAlgOID();
        String signMethod = "";
        String digestMethod = "";
        Provider provider = null;

        // Выбор соответствующего алгоритма подписи
            if(sigAlgOid.equalsIgnoreCase( PKCSObjectIdentifiers.sha1WithRSAEncryption.getId() )) {
                // signMethod =  Constants.MoreAlgorithmsSpecNS + "rsa-sha1"
                // digestMethod = Constants.MoreAlgorithmsSpecNS + "sha1"
                signMethod = XMLSignature.ALGO_ID_SIGNATURE_RSA_SHA1;
                digestMethod = MessageDigestAlgorithm.ALGO_ID_DIGEST_SHA1;
            } else if(sigAlgOid.equalsIgnoreCase(PKCSObjectIdentifiers.sha256WithRSAEncryption.getId()) ) {
                // signMethod = Constants.MoreAlgorithmsSpecNS + "rsa-sha256"
                //signMethod = ALGO_ID_SIGNATURE_RSA_SHA256
                signMethod = Constants.MoreAlgorithmsSpecNS + "rsa-sha256";
                digestMethod = XMLCipherParameters.SHA256;
                provider = new GammaTechProvider();
            }else if(sigAlgOid.equalsIgnoreCase( KZObjectIndentifiers.GOSTR_34_10_2015_SIGNATURE_512.getId() )) {//"1.2.398.3.10.1.1.2.3.2" gost 2012 512
                System.out.println("1.2.398.3.10.1.1.2.3.2 -- "+KZObjectIndentifiers.GOSTR_34_10_2015_SIGNATURE_512.getId());
                signMethod = "urn:ietf:params:xml:ns:pkigovkz:xmlsec:algorithms:gostr34102015-gostr34112015-512";
                digestMethod = "urn:ietf:params:xml:ns:pkigovkz:xmlsec:algorithms:gostr34112015-512";
//                signMethod = "http://www.w3.org/2001/04/xmldsig-more#gostr34102015-gostr34112015-512";
//                digestMethod = "http://www.w3.org/2001/04/xmldsig-more#gostr34112015-512";
            }else if(sigAlgOid.equalsIgnoreCase( KZObjectIndentifiers.GOSTR_34_10_2015_SIGNATURE_256.getId() ))  {//"1.2.398.3.10.1.1.2.3.1" gost 2012 256// не тестировано
                System.out.println("1.2.398.3.10.1.1.2.3.1 -- "+KZObjectIndentifiers.GOSTR_34_10_2015_SIGNATURE_256.getId());
                signMethod = Constants.MoreAlgorithmsSpecNS+"gostr34102012-gostr34112012-256";
                digestMethod = Constants.MoreAlgorithmsSpecNS+"gostr34112012-256";
            }else {
                signMethod = Constants.MoreAlgorithmsSpecNS+"gost34310-gost34311";
                digestMethod = Constants.MoreAlgorithmsSpecNS+"gost34311";
            }
        XMLSignature sig = new XMLSignature(doc, "", signMethod,provider);

        String res = "";
        if (doc.getFirstChild() != null) {
            doc.getFirstChild().appendChild(sig.getElement());
            Transforms transforms = new Transforms(doc);
            transforms.addTransform("http://www.w3.org/2000/09/xmldsig#enveloped-signature");
            transforms.addTransform("http://www.w3.org/TR/2001/REC-xml-c14n-20010315#WithComments");
            sig.addDocument("", transforms, digestMethod);
            sig.addKeyInfo((X509Certificate) cert);
            sig.sign(privKey);
            StringWriter os = new StringWriter();
            TransformerFactory tf = TransformerFactory.newInstance();
            Transformer trans = tf.newTransformer();
            trans.transform(new DOMSource(doc), new StreamResult(os));
            os.flush();
            res = os.toString();
            os.close();
        }
        return parseXml(res);
    }

    public static void testSign(String profile, String pass, String filename, String signFilename) {
        try {
            // Формируем класс хранилища ключей, будут доступны все профайлы криптопровайдера.
            KeyStore store = loadKeyStore(profile,pass);
            //Получение списка ключей
            Enumeration en = store.aliases();
            StoreObjectParam alias = getSignAlias(store);
            if(alias == null)
                throw new Exception("Sign key not found");
            // Получение закрытого ключа по DN имени сертификата
//            PrivateKey prvKey = getPrivateKey("C=KZ, O=AutoTest, CN=Test1521", store, "");
            PrivateKey prvKey = (PrivateKey)store.getKey(alias.getSn(), pass.toCharArray()) ;
            if (prvKey != null) {
                X509Certificate cert = alias.getX509certificate();
                if (cert != null) {
                    DocumentBuilderFactory dbf = DocumentBuilderFactory.newInstance();
                    dbf.setNamespaceAware(true);
                    DocumentBuilder builder = dbf.newDocumentBuilder();
                    InputSource source = null;
                    source = new InputSource(new FileInputStream(filename));
                    // Подписываем XML документ
                    Document sigDoc = signXML(builder.parse(source), cert, prvKey);
                    // Проверяем подпись XML документа
                    if (!validateXML(sigDoc))
                        throw new Exception("Подпись не прошла проверку");
                    else{
                        StringWriter os = new StringWriter();
                        Transformer trans  = TransformerFactory.newInstance().newTransformer();
                        trans.transform(new DOMSource(sigDoc), new StreamResult(os));
                        os.flush();
                        String res = os.toString().replaceAll("&#13;","");
                        os.close();
                        FileWriter fs = new FileWriter(signFilename);
                        fs.write(res);
                        fs.close();
                    }

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

    public static boolean validateXML(Document doc) throws Exception {
        String XMLNS="http://www.w3.org/2000/09/xmldsig#";
        NodeList nl = doc.getElementsByTagNameNS(XMLNS, "Signature");
        System.out.println("sig "+ nl.item(0).toString());
        if(nl.getLength() > 0 ) {
            try {
                Element sigElement = (Element) nl.item(0);
                XMLSignature signature = new XMLSignature(sigElement, "");
                KeyInfo ki = signature.getKeyInfo();
                X509Certificate certKey = ki.getX509Certificate();
                boolean result = false;
                if (certKey != null)
                    result = signature.checkSignatureValue(certKey);

                else {
                    PublicKey pk = ki.getPublicKey();
                    if (pk != null)
                        result = signature.checkSignatureValue(pk);
                    else
                        throw new Exception("Нет информации об открытом ключе. Проверка невозможна.");
                }
                return result;
            }catch (Exception ex){
                ex.printStackTrace();
            }
        }
        return false;
    }

    /**
     *
     * @param keyName   - имя файла ключей без расширения
     * @param path      - путь к файлу с ключами
     * @param pass      - пароль или "" если пароля нет
     * @param ext       - расширение имени файла ключей (bin , p12)
     * @return
     */
    public static String createProfile(String keyName, String path, String pass, String ext) {
        String profile = "";
        Number hProvLocal = TumarCspFunctions.cpAcquireContext("", LibraryWrapper.CRYPT_VERIFYCONTEXT);
        profile = TumarCspFunctions.cpCreateUrl(keyName, "file", keyName, pass, path, ext, 0xA045, 0xAA3A, hProvLocal);
        TumarCspFunctions.cpReleaseContext(hProvLocal, 0);
        return profile;
    }

}
