/*
 * Copyright 2002-2004 The Apache Software Foundation.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * imitations under the License.
 */

package test;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.KeyPairGenerator;
import java.security.KeyPair;
import java.util.Arrays;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;

import junit.framework.Test;
import junit.framework.TestCase;
import junit.framework.TestSuite;

/**
 * Unit tests to test the RSA cipher.
 * 
 * The class generates cipher data with the openSSL RSA cipher functions and
 * compares the results with cipher data generated by the BouncyCastle
 * implementation and vices versa.
 * 
 * Created on 01.10.2005
 * 
 * @author Werner Dittmann (werner@apache.org)
 * 
 */

public class OpenSSLTestRSA extends TestCase {

    KeyPair keyPair = null;

    Cipher cipherEncJuicePub = null;

    Cipher cipherEncBCPub = null;

    Cipher cipherEncJuicePriv = null;

    Cipher cipherEncBCPriv = null;

    Cipher cipherDecJuicePub = null;

    Cipher cipherDecBCPub = null;

    Cipher cipherDecJuicePriv = null;

    Cipher cipherDecBCPriv = null;

    byte[] plainData = null;

    byte[] cipherData_1 = null;

    byte[] cipherData_2 = null;

    byte[] newPlainData_1 = null;

    byte[] newPlainData_2 = null;
    
    public OpenSSLTestRSA(String name) {
        super(name);
        try {
            Class<?> c = Class
                    .forName("org.apache.security.juice.provider.JuiCEProviderOpenSSL");
            java.security.Security.addProvider((java.security.Provider) c
                    .newInstance());
            c = Class
                    .forName("org.bouncycastle.jce.provider.BouncyCastleProvider");
            java.security.Security.addProvider((java.security.Provider) c
                    .newInstance());
        } catch (Throwable t) {
        }

        try {
            KeyPairGenerator keyGen = KeyPairGenerator.getInstance("RSA",
                    "BC");
            keyGen.initialize(1024);
            keyPair = keyGen.generateKeyPair();

        } catch (NoSuchAlgorithmException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        } catch (NoSuchProviderException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        }
    }

    public void setup(boolean shortData) {
        if (shortData) {
            plainData = new byte[] { (byte) 0x54, (byte) 0x85, (byte) 0x9b,
                    (byte) 0x34, (byte) 0x2c, (byte) 0x49, (byte) 0xea,
                    (byte) 0x2a };
        } 
        else {
            plainData = new byte[128];
            for (int i = 0; i < plainData.length; i++) {
                plainData[i] = (byte) (i + 1);
            }
        }

    }

    private void setCiphers(String cipherName) {

        try {
            /*
             * initialize ciphers: use the openSSL implementation and the
             * standard BC implementation (as reference)
             */
            cipherEncBCPub = Cipher.getInstance(cipherName, "BC");
            cipherDecBCPub = Cipher.getInstance(cipherName, "BC");

            cipherEncBCPriv = Cipher.getInstance(cipherName, "BC");
            cipherDecBCPriv = Cipher.getInstance(cipherName, "BC");

            cipherEncJuicePub = Cipher.getInstance(cipherName, "JuiCE");
            assertNotNull("Cipher RSA not found", cipherEncJuicePub);
            cipherDecJuicePub = Cipher.getInstance(cipherName, "JuiCE");

            cipherEncJuicePriv = Cipher.getInstance(cipherName, "JuiCE");
            cipherDecJuicePriv = Cipher.getInstance(cipherName, "JuiCE");

            
            /*
             * Check if cipher providers are correct, i.e the correct provider
             * was chosen.
             */
            assertEquals("Wrong class of Juice provider", cipherEncJuicePub
                    .getProvider().getClass().getName(),
                    "org.apache.security.juice.provider.JuiCEProviderOpenSSL");

            assertEquals("Wrong class of BC provider", cipherEncBCPub
                    .getProvider().getClass().getName(),
                    "org.bouncycastle.jce.provider.BouncyCastleProvider");
            /*
             * Initialize the Ciphers
             */
            try {
                cipherEncJuicePub
                        .init(Cipher.ENCRYPT_MODE, keyPair.getPublic());
                cipherEncJuicePriv.init(Cipher.ENCRYPT_MODE, keyPair
                        .getPrivate());
                cipherDecJuicePub
                        .init(Cipher.DECRYPT_MODE, keyPair.getPublic());
                cipherDecJuicePriv.init(Cipher.DECRYPT_MODE, keyPair
                        .getPrivate());

                cipherEncBCPub.init(Cipher.ENCRYPT_MODE, keyPair.getPublic());
                cipherEncBCPriv.init(Cipher.ENCRYPT_MODE, keyPair.getPrivate());
                cipherDecBCPub.init(Cipher.DECRYPT_MODE, keyPair.getPublic());
                cipherDecBCPriv.init(Cipher.DECRYPT_MODE, keyPair.getPrivate());
            } catch (InvalidKeyException e) {
                // TODO Auto-generated catch block
                e.printStackTrace();
            }

        } catch (NoSuchAlgorithmException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        } catch (NoSuchProviderException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        } catch (NoSuchPaddingException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        }
    }

    byte[] encrypt(Cipher enc) {
        byte[] out = null;
        try {
//            System.out.println("Plain data length: " + plainData.length);
            out = enc.doFinal(plainData);
        } catch (IllegalStateException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        } catch (IllegalBlockSizeException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        } catch (BadPaddingException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        }
        return out;
    }

    byte[] decrypt(Cipher dec, byte[] cdata) {
        byte[] out = null;
        try {
            out = dec.doFinal(cdata);
        } catch (IllegalStateException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        } catch (IllegalBlockSizeException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        } catch (BadPaddingException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        }
        return out;
    }

    void performTest(Cipher enc1, Cipher enc2, Cipher dec1, Cipher dec2) {
        cipherData_1 = encrypt(enc1);
        cipherData_2 = encrypt(enc2);
        assertTrue("Cipher data mismatch", Arrays.equals(cipherData_1,
                cipherData_2));

        newPlainData_1 = decrypt(dec1, cipherData_1);
        newPlainData_2 = decrypt(dec2, cipherData_2);

        assertTrue("New Plain data mismatch", Arrays.equals(newPlainData_1,
                newPlainData_2));
    }

    void performTestPadding(Cipher enc1, Cipher dec1) {
        cipherData_1 = encrypt(enc1);

        newPlainData_1 = decrypt(dec1, cipherData_1);

        assertTrue("New Plain data mismatch", Arrays.equals(newPlainData_1,
                plainData));
    }
    /**
     * Test public encryption using openSSL, private decryption using BC.
     * 
     */
    public void testSSL_BC_pub() {
        setCiphers("RSA");        
        setup(false);
        performTest(cipherEncJuicePub, cipherEncBCPub, cipherDecBCPriv,
                cipherDecJuicePriv);
    }

    public void testSSL_BC_priv() {
        setCiphers("RSA");
        setup(false);
        performTest(cipherEncJuicePriv, cipherEncBCPriv, cipherDecBCPub,
                cipherDecJuicePub);
    }

    public void testSSL_BC_PKCS1_pub() {
        setCiphers("RSA/NONE/PKCS1PADDING");        
        setup(true);
        performTestPadding(cipherEncJuicePub, cipherDecJuicePriv);
        performTestPadding(cipherEncJuicePub, cipherDecBCPriv);
        performTestPadding(cipherEncBCPub, cipherDecJuicePriv);
    }

    public void testSSL_BC_PKCS1_priv() {
        setCiphers("RSA/NONE/PKCS1PADDING");        
        setup(true);
        performTestPadding(cipherEncJuicePriv, cipherDecJuicePub);
        performTestPadding(cipherEncJuicePriv, cipherDecBCPub);
        performTestPadding(cipherEncBCPriv, cipherDecJuicePub);
    }
    
    public void testSSL_BC_OAEP_pub() {
        setCiphers("RSA/NONE/OAEPPADDING");
        setup(true);
        performTestPadding(cipherEncJuicePub, cipherDecJuicePriv);
        performTestPadding(cipherEncJuicePub, cipherDecBCPriv);
        performTestPadding(cipherEncBCPub, cipherDecJuicePriv);
    }

    /*
     * OAEP Padding with private encryption/public decryption is not
     * supported by openSSL
     */
//    public void testSSL_BC_OAEP_priv() {
//        setCiphers("RSA/OAEP");
//        setup(true);
//        performTestPadding(cipherEncJuicePriv, cipherDecJuicePub);
//        performTestPadding(cipherEncJuicePriv, cipherDecBCPub);
//        performTestPadding(cipherEncBCPriv, cipherDecJuicePub);
//    }
    /*
     * If nothing spefied, ECB is default
     */

    /**
     * JUnit suite <p/>
     * 
     * @return a junit test suite
     */
    public static Test suite() {
        return new TestSuite(OpenSSLTestRSA.class);
    }

    /**
     * @param args
     */
    public static void main(String[] args) {
        junit.textui.TestRunner.run(suite());
    }

}
