package ru.rutoken.samples.sunPkcs11Wrapper;

import ru.rutoken.Pkcs11.*;
import sun.security.pkcs11.wrapper.*;
import sun.security.pkcs11.wrapper.CK_ATTRIBUTE;
import sun.security.pkcs11.wrapper.CK_MECHANISM;

public class GOST3410 {

    static PKCS11 pkcs11Wrapper;

    private final static char[] USER_PIN = {'1', '2', '3', '4', '5', '6', '7', '8'};
    private final static byte[] ATTR_CRYPTO_PRO_A = {
            0x06, 0x07, 0x2A, (byte)0x85, 0x03, 0x02, 0x02, 0x23, 0x01};
    private final static byte[] ATTR_CRYPTO_PRO_HASH1 = {
            0x06, 0x07, 0x2A, (byte)0x85, 0x03, 0x02, 0x02, 0x1e, 0x01};

    static byte[] message = {
            0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
            0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
            0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
            0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07 };

    static CK_ATTRIBUTE[] privGostKeyAttr = {
            new CK_ATTRIBUTE(PKCS11Constants.CKA_CLASS, PKCS11Constants.CKO_PRIVATE_KEY),
            new CK_ATTRIBUTE(PKCS11Constants.CKA_LABEL,
                    new String("Sample GOST R 34.10-2001 Private Key (Aktiv Co.)").toCharArray()),
            new CK_ATTRIBUTE(PKCS11Constants.CKA_ID,
                    new String("Sample GOST R 34.10-2001 Keypair (Aktiv Co.)").toCharArray()),
            new CK_ATTRIBUTE(PKCS11Constants.CKA_KEY_TYPE, RtPkcs11Constants.CKK_GOSTR3410.longValue()),
            new CK_ATTRIBUTE(PKCS11Constants.CKA_TOKEN, true),
            new CK_ATTRIBUTE(PKCS11Constants.CKA_PRIVATE, true),
            new CK_ATTRIBUTE(RtPkcs11Constants.CKA_GOSTR3410_PARAMS.longValue(), ATTR_CRYPTO_PRO_A)
    };

    static CK_ATTRIBUTE[] pubGostKeyAttr = {
            new CK_ATTRIBUTE(PKCS11Constants.CKA_CLASS, PKCS11Constants.CKO_PUBLIC_KEY),
            new CK_ATTRIBUTE(PKCS11Constants.CKA_LABEL,
                    new String("Sample GOST R 34.10-2001 Public Key (Aktiv Co.)").toCharArray()),
            new CK_ATTRIBUTE(PKCS11Constants.CKA_ID,
                    new String("Sample GOST R 34.10-2001 Keypair (Aktiv Co.)").toCharArray()),
            new CK_ATTRIBUTE(PKCS11Constants.CKA_KEY_TYPE, RtPkcs11Constants.CKK_GOSTR3410.longValue()),
            new CK_ATTRIBUTE(PKCS11Constants.CKA_TOKEN, true),
            new CK_ATTRIBUTE(PKCS11Constants.CKA_PRIVATE, false),
            new CK_ATTRIBUTE(RtPkcs11Constants.CKA_GOSTR3410_PARAMS.longValue(), ATTR_CRYPTO_PRO_A),
            new CK_ATTRIBUTE(RtPkcs11Constants.CKA_GOSTR3411_PARAMS.longValue(), ATTR_CRYPTO_PRO_HASH1)
    };

    static CK_ATTRIBUTE[] signGostKeyAttr = {
            new CK_ATTRIBUTE(PKCS11Constants.CKA_CLASS, PKCS11Constants.CKO_PRIVATE_KEY),
            new CK_ATTRIBUTE(PKCS11Constants.CKA_LABEL,
                    new String("Sample GOST R 34.10-2001 Private Key (Aktiv Co.)").toCharArray())
    };

    static CK_ATTRIBUTE[] verifyGostKeyAttr = {
            new CK_ATTRIBUTE(PKCS11Constants.CKA_CLASS, PKCS11Constants.CKO_PUBLIC_KEY),
            new CK_ATTRIBUTE(PKCS11Constants.CKA_LABEL,
                    new String("Sample GOST R 34.10-2001 Public Key (Aktiv Co.)").toCharArray())
    };

    static CK_ATTRIBUTE[] value = {
            new CK_ATTRIBUTE(PKCS11Constants.CKA_VALUE)
    };

    static long[] FindObject(long session, CK_ATTRIBUTE[] Template) throws PKCS11Exception{
        try {
            long[] foundObjects;
            pkcs11Wrapper.C_FindObjectsInit(session, Template);
            foundObjects = pkcs11Wrapper.C_FindObjects(session, 100);
            pkcs11Wrapper.C_FindObjectsFinal(session);
            System.out.println("Search has been completed. Objects found: " + foundObjects.length);
            return foundObjects;
        }
        catch(PKCS11Exception e){
            throw e;
        }
    };

    static byte[] HashData(long session, byte[] message) throws PKCS11Exception{
        try {
            System.out.println("Data hashing");

            byte[] digestedMessage = new byte[message.length];
            pkcs11Wrapper.C_DigestSingle(session, new CK_MECHANISM(RtPkcs11Constants.CKM_GOSTR3411.longValue()),
                    message, 0, message.length,
                    digestedMessage, 0, digestedMessage.length);

            System.out.println("Hashed data:");
            for (int i = 0; i < digestedMessage.length; ++i)
            {
                System.out.printf(" %02X", digestedMessage[i]);
                if ((i + 1) % 8 == 0)
                    System.out.println();
            }
            System.out.println("Hashing has been completed.");
            return digestedMessage;
        }
        catch(PKCS11Exception e){
            throw e;
        }
    };

    public static String bytesToHex(byte[] bytes) {
        char[] hexArray = "0123456789ABCDEF".toCharArray();
        char[] hexChars = new char[bytes.length * 2];
        for ( int j = 0; j < bytes.length; ++j ) {
            int v = bytes[j] & 0xFF;
            hexChars[j * 2] = hexArray[v >>> 4];
            hexChars[j * 2 + 1] = hexArray[v & 0x0F];
        }
        return new String(hexChars);
    }

    public static void main(String[] args) {
        long session = -1;

        try {
            System.out.println("Example of working with GOSTR3410 algorithm using rtpkcs11ecp " +
                    "via sun.security.pkcs11.wrapper");

            System.out.println("Library initialization and acquiring of function list");
            pkcs11Wrapper = PKCS11.getInstance("rtpkcs11ecp", "C_GetFunctionList", new sun.security.pkcs11.wrapper.CK_C_INITIALIZE_ARGS(), false);

            System.out.println("Acquiring list of slots with connected tokens");
            long[] slotList = pkcs11Wrapper.C_GetSlotList(true);
            if (slotList.length == 0) {
                throw new Exception("No Rutoken is available!");
            }

            System.out.println("Opening session");
            session = pkcs11Wrapper.C_OpenSession( slotList[0],
                    PKCS11Constants.CKF_SERIAL_SESSION | PKCS11Constants.CKF_RW_SESSION, null, null);

            System.out.println("Logging in as user");
            pkcs11Wrapper.C_Login(session, PKCS11Constants.CKU_USER, USER_PIN);

            System.out.println("Generating GOST R 34.10-2001 key pair");
            pkcs11Wrapper.C_GenerateKeyPair(session,
                    new CK_MECHANISM(RtPkcs11Constants.CKM_GOSTR3410_KEY_PAIR_GEN.longValue()), pubGostKeyAttr, privGostKeyAttr);

            System.out.println("Acquiring key for signing");
            long[] foundObjects = FindObject(session, signGostKeyAttr);
            if (foundObjects.length == 0)
            {
                throw new Exception("No signature key found!");
            }

            byte[] digest = HashData(session, message);

            System.out.println("Data signing");
            pkcs11Wrapper.C_SignInit(session, new CK_MECHANISM(RtPkcs11Constants.CKM_GOSTR3410.longValue()), foundObjects[0]);
            byte[] signature = pkcs11Wrapper.C_Sign(session, digest);

            System.out.println("Signed data:");
            for (int i = 0; i < signature.length; ++i)
            {
                System.out.printf(" %02X", signature[i]);
                if ((i + 1) % 8 == 0)
                    System.out.println();
            }
            System.out.println("Data has been signed successfully.");

            System.out.println("Acquiring key for verification");
            foundObjects = FindObject(session, verifyGostKeyAttr);
            if (foundObjects.length == 0)
            {
                throw new Exception("No verification key found!\n");
            }

            pkcs11Wrapper.C_GetAttributeValue(session, foundObjects[0], value);
            byte[] pubKeyValue = value[0].getByteArray();
            System.out.println("Public key value: " + bytesToHex(pubKeyValue));

            System.out.println("Verifying signature ");
            pkcs11Wrapper.C_VerifyInit(session, new CK_MECHANISM(RtPkcs11Constants.CKM_GOSTR3410.longValue()), foundObjects[0]);
            pkcs11Wrapper.C_Verify(session, digest, signature);
            System.out.println("Verifying has been completed successfully");

            System.out.println("Searching objects on token");
            CK_ATTRIBUTE[][] Template = {signGostKeyAttr, verifyGostKeyAttr};
            for (int j = 0; j < Template.length; ++j) {
                foundObjects = FindObject(session, Template[j]);
                for (int i = 0; i < foundObjects.length; ++i) {
                    System.out.println("Destroying object " + (i+1));
                    pkcs11Wrapper.C_DestroyObject(session, foundObjects[i]);
                }
            }

            System.out.println("Logging out");
            pkcs11Wrapper.C_Logout(session);

            System.out.println("Closing session");
            pkcs11Wrapper.C_CloseSession(session);

            session = -1;

            System.out.println("Test has been completed successfully");
        }
        catch(Exception e){
            System.out.println("Some error occurred. Error code: " + e.getMessage());
        }
        try {
            pkcs11Wrapper.C_Logout(session);
        }
        catch(Exception e){}
        try {
            pkcs11Wrapper.C_CloseSession(session);
        }
        catch(Exception e){}
        try {
            pkcs11Wrapper.C_Finalize(null);
        }
        catch(Exception e){}
    }
}
