package ru.rutoken.samples.pkcs11Wrapper;

import com.sun.jna.Memory;
import com.sun.jna.NativeLong;
import com.sun.jna.ptr.NativeLongByReference;
import ru.rutoken.Pkcs11.*;

public class ExtendedFunctions {
    private final static byte[] USER_PIN = {'1', '2', '3', '4', '5', '6', '7', '8'};
    private final static byte[] NEW_PIN = {'1', '2', '3', '4', '5', '6', '7', '8', '9'};
    private final static byte[] SO_PIN = {'8', '7', '6', '5', '4', '3', '2', '1'};
    private final static String TOKEN_LABEL = "It is my Rutoken label sample ";
    private final static String TOKEN_LONG_LABEL = "!!!Sample of Rutoken long-long-long-long-long label!!!";
    private final static int MAX_ADMIN_RETRY_COUNT = 10;
    private final static int MAX_USER_RETRY_COUNT = 10;

    public static void main(String[] args) {
        NativeLong hSession = Pkcs11Constants.CK_INVALID_HANDLE;
        try {
            System.out.println("Example of rtpkcs11ecp extended functions usage via JNA");

            System.out.println("Library initialization and acquiring of function list");
            NativeLong rv = RtPkcs11Library.getInstance().C_Initialize(null);
            if (!rv.equals(Pkcs11Constants.CKR_OK))
            {
                throw new Exception("Failed to initialize library, error code:" + Long.toHexString(rv.longValue()));
            }

            System.out.println("Acquiring list of slots with connected tokens");
            NativeLongByReference slotsCount = new NativeLongByReference();
            rv = RtPkcs11Library.getInstance().C_GetSlotList(true, null, slotsCount);
            if (!rv.equals(Pkcs11Constants.CKR_OK))
            {
                throw new Exception("C_GetSlotList failed, error code: " + Long.toHexString(rv.longValue()));
            }
            if (slotsCount.getValue().intValue() == 0)
            {
                throw new Exception("No Rutoken is available!");
            }
            NativeLong[] pSlotList = new NativeLong[slotsCount.getValue().intValue()];
            rv = RtPkcs11Library.getInstance().C_GetSlotList(true, pSlotList, slotsCount);
            if (!rv.equals(Pkcs11Constants.CKR_OK))
            {
                throw new Exception("C_GetSlotList failed, error code: " + Long.toHexString(rv.longValue()));
            }

            System.out.println("Acquiring token extended info");
            CK_TOKEN_INFO_EXTENDED tokenInfoEx = new CK_TOKEN_INFO_EXTENDED();
            tokenInfoEx.ulSizeofThisStructure = new NativeLong(tokenInfoEx.size());
            rv = RtPkcs11Ex.getInstance().C_EX_GetTokenInfoExtended(pSlotList[0], tokenInfoEx);
            if (!rv.equals(Pkcs11Constants.CKR_OK))
            {
                throw new Exception("C_EX_GetTokenInfoExtended failed, error code: " + Long.toHexString(rv.longValue()));
            }

            System.out.println("Token info:");
            System.out.printf(" Token type: \t\t\t\t0x%08X ", tokenInfoEx.ulTokenType.longValue());
            if (tokenInfoEx.ulTokenType.equals(RtPkcs11Constants.TOKEN_TYPE_RUTOKEN_ECP))
                System.out.println("(Rutoken ECP)");
            if (tokenInfoEx.ulTokenType.equals(RtPkcs11Constants.TOKEN_TYPE_RUTOKEN_LITE))
                System.out.println("(Rutoken Lite)");
            if (tokenInfoEx.ulTokenType.equals(RtPkcs11Constants.TOKEN_TYPE_RUTOKEN))
                System.out.println("(Rutoken)");
            System.out.printf(" Protocol number: \t\t\t0x%08X \n", tokenInfoEx.ulProtocolNumber.longValue());
            System.out.printf(" Microcode number:  \t\t0x%08X \n", tokenInfoEx.ulMicrocodeNumber.longValue());
            System.out.printf(" Flags: \t\t\t\t\t0x%08X\n", tokenInfoEx.flags.longValue());
            System.out.printf(" Min Admin PIN length: \t\t0x%08X \n", tokenInfoEx.ulMinAdminPinLen.longValue());
            System.out.printf(" Max User PIN length: \t\t0x%08X\n", tokenInfoEx.ulMaxUserPinLen.longValue());
            System.out.printf(" Min User PIN length: \t\t0x%08X\n", tokenInfoEx.ulMinUserPinLen.longValue());
            System.out.printf(" Max Admin login failures: \t0x%08X \n", tokenInfoEx.ulMaxAdminRetryCount.longValue());
            System.out.printf(" Admin login failures: \t\t0x%08X \n", tokenInfoEx.ulAdminRetryCountLeft.longValue());
            System.out.printf(" Max User login failures: \t0x%08X \n", tokenInfoEx.ulMaxUserRetryCount.longValue());
            System.out.printf(" User login failures: \t\t0x%08X \n", tokenInfoEx.ulUserRetryCountLeft.longValue());
            System.out.printf(" Token SN: \t\t\t\t\t0x");
            for (int i = 4; i < tokenInfoEx.serialNumber.length; ++i)
                System.out.printf("%02X", tokenInfoEx.serialNumber[i]);
            System.out.println();
            System.out.printf(" Total memory: \t\t\t\t0x%08X\n", tokenInfoEx.ulTotalMemory.longValue());
            System.out.printf(" Free  memory: \t\t\t\t0x%08X\n", tokenInfoEx.ulFreeMemory.longValue());
            System.out.printf(" ATR: \t\t\t\t\t\t");
            for (int i = 0; i < tokenInfoEx.ulATRLen.longValue(); ++i)
                System.out.printf("%02X", tokenInfoEx.ATR[i]);

            System.out.println("\nToken initialization");
            CK_RUTOKEN_INIT_PARAM initInfo = new CK_RUTOKEN_INIT_PARAM();
            initInfo.UseRepairMode = new NativeLong(0);
            initInfo.pNewAdminPin = new Memory(SO_PIN.length);
            initInfo.pNewAdminPin.write(0, SO_PIN, 0, SO_PIN.length);
            initInfo.ulNewAdminPinLen = new NativeLong(SO_PIN.length);
            initInfo.pNewUserPin = new Memory(NEW_PIN.length);
            initInfo.pNewUserPin.write(0, NEW_PIN, 0, NEW_PIN.length);
            initInfo.ulNewUserPinLen = new NativeLong(NEW_PIN.length);
            initInfo.ulMinAdminPinLen = new NativeLong(6);
            initInfo.ulMinUserPinLen = new NativeLong(6);
            initInfo.ChangeUserPINPolicy = new NativeLong(RtPkcs11Constants.TOKEN_FLAGS_ADMIN_CHANGE_USER_PIN.longValue() |
                                                          RtPkcs11Constants.TOKEN_FLAGS_USER_CHANGE_USER_PIN.longValue());
            initInfo.ulMaxAdminRetryCount = new NativeLong(MAX_ADMIN_RETRY_COUNT);
            initInfo.ulMaxUserRetryCount = new NativeLong(MAX_USER_RETRY_COUNT);
            initInfo.pTokenLabel = new Memory(TOKEN_LABEL.length());
            initInfo.pTokenLabel.write(0, TOKEN_LABEL.getBytes(), 0, TOKEN_LABEL.length());
            initInfo.ulLabelLen = new NativeLong(TOKEN_LABEL.length());
            initInfo.ulSmMode = new NativeLong(0);
            initInfo.ulSizeofThisStructure = new NativeLong(initInfo.size());

            rv = RtPkcs11Ex.getInstance().C_EX_InitToken(pSlotList[0], SO_PIN, new NativeLong(SO_PIN.length), initInfo);
            if (!rv.equals(Pkcs11Constants.CKR_OK))
            {
                throw new Exception("C_EX_InitToken failed, error code: " + Long.toHexString(rv.longValue()));
            }

            System.out.println("Opening session");
            NativeLongByReference phSession = new NativeLongByReference();
            rv = RtPkcs11Library.getInstance().C_OpenSession(pSlotList[0],
                    new NativeLong(Pkcs11Constants.CKF_SERIAL_SESSION.longValue() |
                                   Pkcs11Constants.CKF_RW_SESSION.longValue()),
                    null, null, phSession);
            if (!rv.equals(Pkcs11Constants.CKR_OK))
            {
                throw new Exception("C_OpenSession failed, error code: " + Long.toHexString(rv.longValue()));
            }
            hSession = phSession.getValue();

            System.out.println("Logging in as administrator");
            rv = RtPkcs11Library.getInstance().C_Login(hSession, Pkcs11Constants.CKU_SO, SO_PIN, new NativeLong(SO_PIN.length));
            if (!rv.equals(Pkcs11Constants.CKR_OK))
            {
                throw new Exception("C_Login failed, error code: " + Long.toHexString(rv.longValue()));
            }

            System.out.println("User PIN unlocking");
            rv = RtPkcs11Ex.getInstance().C_EX_UnblockUserPIN(hSession);
            if (!rv.equals(Pkcs11Constants.CKR_OK))
            {
                throw new Exception("C_EX_UnblockUserPIN failed, error code: " + Long.toHexString(rv.longValue()));
            }

            System.out.println("Logging out");
            rv = RtPkcs11Library.getInstance().C_Logout(hSession);
            if (!rv.equals(Pkcs11Constants.CKR_OK))
            {
                throw new Exception("C_Logout failed, error code: " + Long.toHexString(rv.longValue()));
            }

            System.out.println("Logging in as user");
            rv = RtPkcs11Library.getInstance().C_Login(hSession, Pkcs11Constants.CKU_USER,
                                                       NEW_PIN, new NativeLong(NEW_PIN.length));
            if (!rv.equals(Pkcs11Constants.CKR_OK))
            {
                throw new Exception("C_Login failed, error code: " + Long.toHexString(rv.longValue()));
            }

            System.out.println("Setting token name");
            rv = RtPkcs11Ex.getInstance().C_EX_SetTokenName(hSession, TOKEN_LONG_LABEL.getBytes(),
                    new NativeLong(TOKEN_LONG_LABEL.length()));
            if (!rv.equals(Pkcs11Constants.CKR_OK))
            {
                throw new Exception("C_EX_SetTokenName failed, error code: " + Long.toHexString(rv.longValue()));
            }

            System.out.println("User PIN changing");
            rv = RtPkcs11Library.getInstance().C_SetPIN(hSession, NEW_PIN, new NativeLong(NEW_PIN.length),
                                                                  USER_PIN, new NativeLong(USER_PIN.length));
            if (!rv.equals(Pkcs11Constants.CKR_OK))
            {
                throw new Exception("C_SetPIN failed, error code: " + Long.toHexString(rv.longValue()));
            }
        }
        catch(Exception e){
            System.out.println(e.getMessage());
        }
        finally {
            System.out.println("Logging out");
            NativeLong rv = RtPkcs11Library.getInstance().C_Logout(hSession);
            if (!rv.equals(Pkcs11Constants.CKR_OK))
            {
                System.out.println("C_Logout failed, error code: " + Long.toHexString(rv.longValue()));
            }

            System.out.println("Closing session");
            rv = RtPkcs11Library.getInstance().C_CloseSession(hSession);
            if (!rv.equals(Pkcs11Constants.CKR_OK))
            {
                System.out.println("C_CloseSession failed, error code: " + Long.toHexString(rv.longValue()));
            }

            System.out.println("Finalizing PKCS11 library");
            rv = RtPkcs11Library.getInstance().C_Finalize(null);
            if (!rv.equals(Pkcs11Constants.CKR_OK))
            {
                System.out.println("C_Finalize failed, error code: " + Long.toHexString(rv.longValue()));
            }

            System.out.println("Test has been completed.");
        }
    }
}

