/*
 * Copyright (c) 2014, Oracle and/or its affiliates. All rights reserved.
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 * 
 * This code is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License version 2 only, as
 * published by the Free Software Foundation.
 * 
 * This code is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
 * version 2 for more details (a copy is included in the LICENSE file that
 * accompanied this code).
 * 
 * You should have received a copy of the GNU General Public License version
 * 2 along with this work; if not, write to the Free Software Foundation,
 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
 * 
 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
 * or visit www.oracle.com if you need additional information or have any
 * questions.
 */

/**
 * Inter-operability testing between SunJCE and JSafeJCE.
 * javac -Djava.ext.dirs="lib/cryptoj-6.1:jre/lib/ext" TestInteropCryptoJ.java 
 * java -Djava.ext.dirs="lib/cryptoj-6.1:jre/lib/ext" TestInteropCryptoJ
 *
 * @author rhalade
 */
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.Provider;
import java.security.PublicKey;
import java.security.Security;
import java.security.spec.X509EncodedKeySpec;
import javax.crypto.KeyAgreement;
import javax.crypto.ShortBufferException;
import javax.crypto.interfaces.DHPublicKey;
import javax.crypto.spec.DHParameterSpec;

public class TestInteropCryptoJ {

    public static void main(String argv[]) throws Exception {
        
        String sunjce = "SunJCE";
        String jsafejce = "JsafeJCE";
        
        Provider jsafeProvider = new com.rsa.jsafe.provider.JsafeJCE();
        Security.insertProviderAt(jsafeProvider, 1);

        // ==> Alice - SunJCE
        // Create DH key pair
        KeyPairGenerator aliceKpairGen
                = KeyPairGenerator.getInstance("DH", sunjce);
        aliceKpairGen.initialize(1024);
        KeyPair aliceKpair = aliceKpairGen.generateKeyPair();

        // Create and initialize DH KeyAgreement object
        KeyAgreement aliceKeyAgree = KeyAgreement.getInstance("DH", sunjce);
        aliceKeyAgree.init(aliceKpair.getPrivate());

        // Encode public key, and send to Bob
        byte[] alicePubKeyEnc = aliceKpair.getPublic().getEncoded();

        // ==> Bob - JsafeJCE
        // Create public key from encoded data
        KeyFactory bobKeyFac = KeyFactory.getInstance("DH", jsafejce);
        X509EncodedKeySpec x509KeySpec = new X509EncodedKeySpec(
                alicePubKeyEnc);
        PublicKey alicePubKey = bobKeyFac.generatePublic(x509KeySpec);

        // Retrieve DH parameters used
        DHParameterSpec dhParamSpec = ((DHPublicKey) alicePubKey).getParams();

        // Create DH key pair
        KeyPairGenerator bobKpairGen
                = KeyPairGenerator.getInstance("DH", jsafejce);
        bobKpairGen.initialize(dhParamSpec);
        KeyPair bobKpair = bobKpairGen.generateKeyPair();

        // Create and initialize DH KeyAgreement object
        KeyAgreement bobKeyAgree = KeyAgreement.getInstance("DH", jsafejce);
        bobKeyAgree.init(bobKpair.getPrivate());

        // Perform only phase of DH protocol
        bobKeyAgree.doPhase(alicePubKey, true);

        // Encode public key, and send to Alice
        byte[] bobPubKeyEnc = bobKpair.getPublic().getEncoded();

        // ==> Alice - SunJCE
        // Create public key from encoded data
        KeyFactory aliceKeyFac = KeyFactory.getInstance("DH", sunjce);
        x509KeySpec = new X509EncodedKeySpec(bobPubKeyEnc);
        PublicKey bobPubKey = aliceKeyFac.generatePublic(x509KeySpec);

        // Perform only phase of DH protocol
        aliceKeyAgree.doPhase(bobPubKey, true);

        // DH key agreement complete, both have same secret
        byte[] aliceSharedSecret = aliceKeyAgree.generateSecret();
        
        /*
         // This works
         byte[] bobSharedSecret = bobKeyAgree.generateSecret();
         */
        
        // This doesn't work when we try with smaller buffer first
        // As per documentation:
        // http://docs.oracle.com/javase/8/docs/api/javax/crypto/KeyAgreement.html#generateSecret-byte:A-int-
        // "If the sharedSecret buffer is too small to hold the result, 
        // a ShortBufferException is thrown. In this case, this call should be 
        // repeated with a larger output buffer."
        int secretLen = aliceSharedSecret.length;
        byte[] bobSharedSecret = new byte[secretLen];
        try {
            bobKeyAgree.generateSecret(bobSharedSecret, 1);
        } catch (ShortBufferException e) {
            System.out.println("Let's ignore this exception: " + e.getMessage());
        }

        // we repeat the call with larger output buffer
        bobKeyAgree.generateSecret(bobSharedSecret, 0);

        if (!java.util.Arrays.equals(aliceSharedSecret, bobSharedSecret)) {
            throw new RuntimeException("Do not share same secret");
        }
        
        System.out.println("PASS");
    }
}