package kz.gamma;

import com.sun.org.apache.xpath.internal.XPathAPI;
import kz.gamma.asn1.x509.X509Name;
import kz.gamma.jce.provider.GammaTechProvider;
import kz.gamma.tumarcsp.params.StoreObjectParam;
import kz.gamma.xmldsig.JCPXMLDSigInit;
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.XMLUtils;
import org.w3c.dom.Document;
import org.w3c.dom.Element;
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.Date;
import java.util.Enumeration;

/**
 * Created by s_moiseyev
 * Date: 12.10.2009
 * Time: 11:08:34
 */
public class SampleJCECSP_SignatureXMLRSA {

    /**
     * @param args
     */
    public static void main(String[] args) {
        Security.addProvider(new GammaTechProvider());
        JCPXMLDSigInit.init();
        //veryfyXMLTest("C:\\xmlTest.xml");
        profileTest();
        //profileTest("profile://TestFileSystem", "");
    }

    /**
     *
     */
    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=AutoTest, CN=Test1521", store, "");
            if (prvKey != null) {
                // Получение сертификата по DN имени
                Certificate cert = getCertificate("C=KZ, O=AutoTest, CN=Test1521", store, "");
                if (cert != null) {
                    DocumentBuilderFactory dbf = DocumentBuilderFactory.newInstance();
                    dbf.setNamespaceAware(true);
                    DocumentBuilder builder = dbf.newDocumentBuilder();
                    InputSource source = null;
                    source = new InputSource(new FileInputStream("C:\\temp\\config.xml"));
                    // Подписываем XML документ
                    Document sigDoc = signXML(builder.parse(source), cert, prvKey);
                    // Проверяем подпись XML документа
                    if (!validateXML(sigDoc))
                        throw new Exception("Подпись не прошла проверку");
                }
            }
        } 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=GAMMA,CN=GAMMAJCE", store, pass);
            if (prvKey != null) {
                // Получение сертификата по DN имени
                Certificate cert = getCertificate("C=KZ,O=GAMMA,CN=GAMMAJCE", store, "");
                if (cert != null) {
                    DocumentBuilderFactory dbf = DocumentBuilderFactory.newInstance();
                    dbf.setNamespaceAware(true);
                    DocumentBuilder builder = dbf.newDocumentBuilder();
                    InputSource source = null;
                    source = new InputSource(new FileInputStream("C:\\temp\\config.xml"));
                    // Подписываем XML документ
                    Document sigDoc = signXML(builder.parse(source), cert, prvKey);
                    // Проверяем подпись XML документа
                    if (!validateXML(sigDoc))
                        throw new Exception("Подпись не прошла проверку");
                }
            }
        } 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 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 ((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 ((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 ((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.getCertificateChain(tmpSN);
    }

    /**
     * Метод формирования подписи xml документа
     *
     * @param doc
     * @param cert
     * @param privKey
     * @return
     * @throws Exception
     */
    public static Document signXML(Document doc, Certificate cert, PrivateKey privKey)
            throws Exception {
        String signMethod = "http://www.w3.org/2001/04/xmldsig-more#rsa-sha1";
        String digestMethod = "http://www.w3.org/2001/04/xmldsig-more#sha1";
        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);
    }

    /**
     * Метод проверки подписи xml документа
     *
     * @param doc
     * @return
     * @throws Exception
     */
    public static boolean validateXML(Document doc)
            throws Exception {
        // Если уже один раз был объявлен данный метод, то его повторно не нужно объявлять
        //JCPXMLDSigInit.init();
        Element nscontext = XMLUtils.createDSctx(doc, "ds", "http://www.w3.org/2000/09/xmldsig#");
        Element sigElement = (Element) XPathAPI.selectSingleNode(doc, "//ds:Signature[1]", nscontext);
        XMLSignature signature = new XMLSignature(sigElement, "");
        //System.out.println(.toString());
        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;
    }

    /**
     * @param xml
     * @return
     * @throws Exception
     */
    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);
    }

    /**
     * @param reader
     * @return
     * @throws IOException
     */
    public static String readBigString(BufferedReader reader) throws IOException {
        StringBuffer buf = new StringBuffer();
        for (String curr = reader.readLine(); curr != null; curr = reader.readLine()) {
            buf.append((new StringBuilder(String.valueOf(curr))).append("\n").toString());
        }
        return buf.toString();
    }

    /**
     * @param xml
     * @return
     */
    public static String stripXml(String xml) {
        String result = xml.trim();
        return result;
    }

    /**
     * @param fileName
     */
    public static void veryfyXMLTest(String fileName) {
        try {
            if (!validateXML(parse(fileName, true))) {
                System.out.println("Подпись не прошла проверку");
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    /**
     * @param fileName
     * @param decode
     * @return
     * @throws Exception
     */
    public static Document parse(String fileName, boolean decode) throws Exception {
        DocumentBuilderFactory dbf = DocumentBuilderFactory.newInstance();
        dbf.setNamespaceAware(true);
        DocumentBuilder builder = dbf.newDocumentBuilder();
        InputSource source = null;
        if (decode) {
            String input = stripXml(readBigString(new BufferedReader(new InputStreamReader(new FileInputStream(fileName), "UTF-8"))));
            input = input.replaceAll("&amp;", "&");
            input = input.replaceAll("&quot;", "\"");
            input = input.replaceAll("&quote;", "\"");
            input = input.replaceAll("&gt;", ">");
            input = input.replaceAll("&lt;", "<");
            System.out.println((new StringBuilder("Input: ")).append(input).toString());
            source = new InputSource(new StringReader(input));
        } else {
            source = new InputSource(new FileInputStream(fileName));
        }
        return builder.parse(source);
    }
}