import java.security.spec.MGF1ParameterSpec;

import java.security.spec.PSSParameterSpec;

import java.util.Locale;

 *  . "RSASSA-PSS"

    protected MessageDigest messageDigest;

    protected boolean needsReset;

    protected Key privateKey = null;

    protected Key publicKey = null;

    public static final class PSS extends RSASignature {

        private PSSParameterSpec pssParams = null;

        // Workaround: Cannot import raw public key to CNG. This signature

        // will be used for verification if key is not from MSCAPI.

        private Signature fallbackSignature;

        @Override

        protected void engineInitSign(PrivateKey key) throws InvalidKeyException {

            super.engineInitSign(key);

            fallbackSignature = null;

        }

        @Override

        protected void engineInitVerify(PublicKey key) throws InvalidKeyException {

            // This signature accepts only RSAPublicKey

            if ((key instanceof java.security.interfaces.RSAPublicKey) == false) {

                throw new InvalidKeyException("Key type not supported");

            }

            this.privateKey = null;

            if (key instanceof sun.security.mscapi.RSAPublicKey) {

                fallbackSignature = null;

                publicKey = (sun.security.mscapi.RSAPublicKey) key;

            } else {

                if (fallbackSignature == null) {

                    try {

                        fallbackSignature = Signature.getInstance(

                                "RSASSA-PSS", "SunRsaSign");

                    } catch (NoSuchAlgorithmException | NoSuchProviderException e) {

                        throw new InvalidKeyException("Invalid key", e);

                    }

                }

                fallbackSignature.initVerify(key);

                if (pssParams != null) {

                    try {

                        fallbackSignature.setParameter(pssParams);

                    } catch (InvalidAlgorithmParameterException e) {

                        throw new InvalidKeyException("Invalid params", e);

                    }

                }

                publicKey = null;

            }

            resetDigest();

        }

        @Override

        protected void engineUpdate(byte b) throws SignatureException {

            ensureInit();

            if (fallbackSignature != null) {

                fallbackSignature.update(b);

            } else {

                messageDigest.update(b);

            }

            needsReset = true;

        }

        @Override

        protected void engineUpdate(byte[] b, int off, int len) throws SignatureException {

            ensureInit();

            if (fallbackSignature != null) {

                fallbackSignature.update(b, off, len);

            } else {

                messageDigest.update(b, off, len);

            }

            needsReset = true;

        }

        @Override

        protected void engineUpdate(ByteBuffer input) {

            try {

                ensureInit();

            } catch (SignatureException se) {

                // hack for working around API bug

                throw new RuntimeException(se.getMessage());

            }

            if (fallbackSignature != null) {

                try {

                    fallbackSignature.update(input);

                } catch (SignatureException se) {

                    // hack for working around API bug

                    throw new RuntimeException(se.getMessage());

                }

            } else {

                messageDigest.update(input);

            }

            needsReset = true;

        }

        @Override

        protected byte[] engineSign() throws SignatureException {

            ensureInit();

            byte[] hash = getDigestValue();

            return signPssHash(hash, hash.length,

                    pssParams.getSaltLength(),

                    ((MGF1ParameterSpec)

                            pssParams.getMGFParameters()).getDigestAlgorithm(),

                    privateKey.getHCryptProvider(), privateKey.getHCryptKey());

        }

        @Override

        protected boolean engineVerify(byte[] sigBytes) throws SignatureException {

            ensureInit();

            if (fallbackSignature != null) {

                needsReset = false;

                return fallbackSignature.verify(sigBytes);

            } else {

                byte[] hash = getDigestValue();

                return verifyPssSignedHash(

                        hash, hash.length,

                        sigBytes, sigBytes.length,

                        pssParams.getSaltLength(),

                        ((MGF1ParameterSpec)

                                pssParams.getMGFParameters()).getDigestAlgorithm(),

                        publicKey.getHCryptProvider(),

                        publicKey.getHCryptKey()

                );

            }

        }

        @Override

        protected void engineSetParameter(AlgorithmParameterSpec params)

                throws InvalidAlgorithmParameterException {

            if (needsReset) {

                throw new ProviderException

                        ("Cannot set parameters during operations");

            }

            this.pssParams = validateSigParams(params);

            if (fallbackSignature != null) {

                fallbackSignature.setParameter(params);

            }

        }

        @Override

        protected AlgorithmParameters engineGetParameters() {

            if (this.pssParams == null) {

                throw new ProviderException("Missing required PSS parameters");

            }

            try {

                AlgorithmParameters ap =

                        AlgorithmParameters.getInstance("RSASSA-PSS");

                ap.init(this.pssParams);

                return ap;

            } catch (GeneralSecurityException gse) {

                throw new ProviderException(gse.getMessage());

            }

        }

        private void ensureInit() throws SignatureException {

            if (this.privateKey == null && this.publicKey == null

                    && fallbackSignature == null) {

                throw new SignatureException("Missing key");

            }

            if (this.pssParams == null) {

                // Parameters are required for signature verification

                throw new SignatureException

                        ("Parameters required for RSASSA-PSS signatures");

            }

            if (fallbackSignature == null && messageDigest == null) {

                // This could happen if initVerify(softKey), setParameter(),

                // and initSign() were called. No messageDigest. Create it.

                try {

                    messageDigest = MessageDigest

                            .getInstance(pssParams.getDigestAlgorithm());

                } catch (NoSuchAlgorithmException e) {

                    throw new SignatureException(e);

                }

            }

        }

        /**

         * Validate the specified Signature PSS parameters.

         */

        private PSSParameterSpec validateSigParams(AlgorithmParameterSpec p)

                throws InvalidAlgorithmParameterException {

            if (p == null) {

                throw new InvalidAlgorithmParameterException

                        ("Parameters cannot be null");

            }

            if (!(p instanceof PSSParameterSpec)) {

                throw new InvalidAlgorithmParameterException

                        ("parameters must be type PSSParameterSpec");

            }

            // no need to validate again if same as current signature parameters

            PSSParameterSpec params = (PSSParameterSpec) p;

            if (params == this.pssParams) return params;

            // now sanity check the parameter values

            if (!(params.getMGFAlgorithm().equalsIgnoreCase("MGF1"))) {

                throw new InvalidAlgorithmParameterException("Only supports MGF1");

            }

            if (params.getTrailerField() != PSSParameterSpec.TRAILER_FIELD_BC) {

                throw new InvalidAlgorithmParameterException

                        ("Only supports TrailerFieldBC(1)");

            }

            AlgorithmParameterSpec algSpec = params.getMGFParameters();

            if (!(algSpec instanceof MGF1ParameterSpec)) {

                throw new InvalidAlgorithmParameterException

                        ("Only support MGF1ParameterSpec");

            }

            MGF1ParameterSpec mgfSpec = (MGF1ParameterSpec)algSpec;

            String msgHashAlg = params.getDigestAlgorithm()

                    .toLowerCase(Locale.ROOT).replaceAll("-", "");

            if (msgHashAlg.equals("sha")) {

                msgHashAlg = "sha1";

            }

            String mgf1HashAlg = mgfSpec.getDigestAlgorithm()

                    .toLowerCase(Locale.ROOT).replaceAll("-", "");

            if (mgf1HashAlg.equals("sha")) {

                mgf1HashAlg = "sha1";

            }

            if (!mgf1HashAlg.equals(msgHashAlg)) {

                throw new InvalidAlgorithmParameterException

                        ("MGF1 hash must be the same as message hash");

            }

            return params;

        }

        /**

         * Sign hash using CNG API with HCRYPTKEY. Used by RSASSA-PSS.

         */

        private native static byte[] signPssHash(byte[] hash,

                int hashSize, int saltLength, String hashAlgorithm,

                long hCryptProv, long nCryptKey)

                throws SignatureException;

        /**

         * Verify a signed hash using CNG API with HCRYPTKEY. Used by RSASSA-PSS.

         * This method is not used now. See {@link #fallbackSignature}.

         */

        private native static boolean verifyPssSignedHash(byte[] hash, int hashSize,

                byte[] signature, int signatureSize,

                int saltLength, String hashAlgorithm,

                long hCryptProv, long hKey) throws SignatureException;

    }