package ru.rutoken.samples.sunPkcs11Wrapper;

import sun.security.pkcs11.wrapper.*;

public class RSA {

    static PKCS11 pkcs11Wrapper;

    private final static char[] USER_PIN = {'1', '2', '3', '4', '5', '6', '7', '8'};

    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[] privRsaKeyAttr = {
            new CK_ATTRIBUTE(PKCS11Constants.CKA_CLASS, PKCS11Constants.CKO_PRIVATE_KEY),
            new CK_ATTRIBUTE(PKCS11Constants.CKA_LABEL, new String("Sample RSA Private Key (Aktiv Co.)").toCharArray()),
            new CK_ATTRIBUTE(PKCS11Constants.CKA_ID, new String("Sample RSA Keypair ID (Aktiv Co.)").toCharArray()),
            new CK_ATTRIBUTE(PKCS11Constants.CKA_KEY_TYPE, PKCS11Constants.CKK_RSA),
            new CK_ATTRIBUTE(PKCS11Constants.CKA_DECRYPT, true),
            new CK_ATTRIBUTE(PKCS11Constants.CKA_TOKEN, true),
            new CK_ATTRIBUTE(PKCS11Constants.CKA_PRIVATE, true)
    };

    static CK_ATTRIBUTE[] pubRsaKeyAttr = {
            new CK_ATTRIBUTE(PKCS11Constants.CKA_CLASS, PKCS11Constants.CKO_PUBLIC_KEY),
            new CK_ATTRIBUTE(PKCS11Constants.CKA_LABEL, new String("Sample RSA Public Key (Aktiv Co.)").toCharArray()),
            new CK_ATTRIBUTE(PKCS11Constants.CKA_ID, new String("Sample RSA Keypair ID (Aktiv Co.)").toCharArray()),
            new CK_ATTRIBUTE(PKCS11Constants.CKA_KEY_TYPE, PKCS11Constants.CKK_RSA),
            new CK_ATTRIBUTE(PKCS11Constants.CKA_ENCRYPT, true),
            new CK_ATTRIBUTE(PKCS11Constants.CKA_TOKEN, true),
            new CK_ATTRIBUTE(PKCS11Constants.CKA_PRIVATE, false),
            new CK_ATTRIBUTE(PKCS11Constants.CKA_MODULUS_BITS, 512)
    };

    static CK_ATTRIBUTE[] signRsaKeyAttr = {
            new CK_ATTRIBUTE(PKCS11Constants.CKA_CLASS, PKCS11Constants.CKO_PRIVATE_KEY),
            new CK_ATTRIBUTE(PKCS11Constants.CKA_LABEL, new String("Sample RSA Private Key (Aktiv Co.)").toCharArray())
    };
    static CK_ATTRIBUTE[] verifyRsaKeyAttr = {
            new CK_ATTRIBUTE(PKCS11Constants.CKA_CLASS, PKCS11Constants.CKO_PUBLIC_KEY),
            new CK_ATTRIBUTE(PKCS11Constants.CKA_LABEL, new String("Sample RSA Public Key (Aktiv Co.)").toCharArray())
    };

    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(PKCS11Constants.CKM_SHA_1),
                    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 void main(String[] args) {
        long session = -1;

        try {
            System.out.println("Example of working with RSA 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 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 RSA key pair");
            pkcs11Wrapper.C_GenerateKeyPair(
                    session, new CK_MECHANISM(PKCS11Constants.CKM_RSA_PKCS_KEY_PAIR_GEN), pubRsaKeyAttr, privRsaKeyAttr);

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

            System.out.println("Data encrypting");
            pkcs11Wrapper.C_EncryptInit(session, new CK_MECHANISM(PKCS11Constants.CKM_RSA_PKCS), foundObjects[0]);
            byte[] encryptedMessage = new byte[message.length*2];
            pkcs11Wrapper.C_Encrypt(session, message, 0, message.length, encryptedMessage, 0, encryptedMessage.length);

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

            System.out.println("Acquiring private key");
            foundObjects = FindObject(session, signRsaKeyAttr);
            if (foundObjects.length == 0)
            {
                throw new Exception("No private key found!");
            }

            System.out.println("Data decrypting");
            pkcs11Wrapper.C_DecryptInit(session, new CK_MECHANISM(PKCS11Constants.CKM_RSA_PKCS), foundObjects[0]);
            byte[] decryptedMessage = new byte[message.length];
            pkcs11Wrapper.C_Decrypt(session, encryptedMessage, 0, encryptedMessage.length, decryptedMessage, 0, decryptedMessage.length);

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

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

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

            System.out.println("Signing data");
            pkcs11Wrapper.C_SignInit(session, new CK_MECHANISM(PKCS11Constants.CKM_RSA_PKCS), 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 verifying");
            foundObjects = FindObject(session, verifyRsaKeyAttr);
            if (foundObjects.length == 0)
            {
                throw new Exception("No verification key found!");
            }

            System.out.println("Verifying data");
            pkcs11Wrapper.C_VerifyInit(session, new CK_MECHANISM(PKCS11Constants.CKM_RSA_PKCS), 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 = {signRsaKeyAttr, verifyRsaKeyAttr};
            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){}
    }
}
