package org.bouncycastle.cms.test;

import java.security.KeyPair;
import java.security.cert.X509Certificate;
import java.util.Arrays;
import java.util.Collection;
import java.util.Iterator;

import javax.crypto.SecretKey;

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

import org.bouncycastle.cms.CMSEnvelopedData;
import org.bouncycastle.cms.CMSEnvelopedDataGenerator;
import org.bouncycastle.cms.CMSProcessableByteArray;
import org.bouncycastle.cms.RecipientInformation;
import org.bouncycastle.cms.RecipientInformationStore;

public class EnvelopedDataTest extends TestCase {

	/*
	 *
	 *  VARIABLES
	 *
	 */

	public boolean DEBUG = true;

	/*
	 *
	 *  INFRASTRUCTURE
	 *
	 */

	public EnvelopedDataTest(String name) {
		super(name);
	}

	public static void main(String args[]) {
		junit.textui.TestRunner.run(EnvelopedDataTest.class);
	}

	public static Test suite() {
		return new TestSuite(EnvelopedDataTest.class);
	}

	public void log(Exception _ex) {
		if(DEBUG) {
			_ex.printStackTrace();
		}
	}

	public void log(String _msg) {
		if(DEBUG) {
			System.out.println(_msg);
		}
	}

	public void setUp() {

	}

	public void tearDown() {

	}

	/*
	 *
	 *  TESTS
	 *
	 */

	public void testKeyTrans()
    {
		try {
			byte[]          data     = "WallaWallaWashington".getBytes();

			String          _signDN   = "O=Bouncy Castle, C=AU";
			KeyPair         _signKP   = CMSTestUtil.makeKeyPair();  
			X509Certificate _signCert = CMSTestUtil.makeCertificate(_signKP, _signDN, _signKP, _signDN);

			String          _origDN   = "CN=Bob, OU=Sales, O=Bouncy Castle, C=AU";
			KeyPair         _origKP   = CMSTestUtil.makeKeyPair();
			X509Certificate _origCert = CMSTestUtil.makeCertificate(_origKP, _origDN, _signKP, _signDN);

			String          _reciDN   = "CN=Doug, OU=Sales, O=Bouncy Castle, C=AU";
			KeyPair         _reciKP   = CMSTestUtil.makeKeyPair();
			X509Certificate _reciCert = CMSTestUtil.makeCertificate(_reciKP, _reciDN, _signKP, _signDN);
			
			CMSEnvelopedDataGenerator edGen = new CMSEnvelopedDataGenerator();

			edGen.addKeyTransRecipient(_reciCert);

			CMSEnvelopedData ed = edGen.generate(
                                    new CMSProcessableByteArray(data),
                                    CMSEnvelopedDataGenerator.DES_EDE3_CBC, "BC");

            RecipientInformationStore  recipients = ed.getRecipientInfos();

            Collection  c = recipients.getRecipients();
            Iterator    it = c.iterator();

            while (it.hasNext())
            {
                RecipientInformation   recipient = (RecipientInformation)it.next();

                byte[] recData = recipient.getContent(_reciKP.getPrivate(), "BC");

                assertEquals(true, Arrays.equals(data, recData));
            }
		}
		catch(Exception ex) {
			log(ex);
			fail();
		}
	}

	public void testKeyTransAES128()
    {
		try {
			byte[]          data     = "WallaWallaWashington".getBytes();

			String          _signDN   = "O=Bouncy Castle, C=AU";
			KeyPair         _signKP   = CMSTestUtil.makeKeyPair();  
			X509Certificate _signCert = CMSTestUtil.makeCertificate(_signKP, _signDN, _signKP, _signDN);

			String          _origDN   = "CN=Bob, OU=Sales, O=Bouncy Castle, C=AU";
			KeyPair         _origKP   = CMSTestUtil.makeKeyPair();
			X509Certificate _origCert = CMSTestUtil.makeCertificate(_origKP, _origDN, _signKP, _signDN);

			String          _reciDN   = "CN=Doug, OU=Sales, O=Bouncy Castle, C=AU";
			KeyPair         _reciKP   = CMSTestUtil.makeKeyPair();
			X509Certificate _reciCert = CMSTestUtil.makeCertificate(_reciKP, _reciDN, _signKP, _signDN);
			
			CMSEnvelopedDataGenerator edGen = new CMSEnvelopedDataGenerator();

			edGen.addKeyTransRecipient(_reciCert);

			CMSEnvelopedData ed = edGen.generate(
                                    new CMSProcessableByteArray(data),
                                    CMSEnvelopedDataGenerator.AES128_CBC, "BC");

            RecipientInformationStore  recipients = ed.getRecipientInfos();

            Collection  c = recipients.getRecipients();
            Iterator    it = c.iterator();

            while (it.hasNext())
            {
                RecipientInformation   recipient = (RecipientInformation)it.next();

                byte[] recData = recipient.getContent(_reciKP.getPrivate(), "BC");

                assertEquals(true, Arrays.equals(data, recData));
            }
		}
		catch(Exception ex) {
			log(ex);
			fail();
		}
	}

	public void testKeyTransAES192()
    {
		try {
			byte[]          data     = "WallaWallaWashington".getBytes();

			String          _signDN   = "O=Bouncy Castle, C=AU";
			KeyPair         _signKP   = CMSTestUtil.makeKeyPair();  
			X509Certificate _signCert = CMSTestUtil.makeCertificate(_signKP, _signDN, _signKP, _signDN);

			String          _reciDN   = "CN=Doug, OU=Sales, O=Bouncy Castle, C=AU";
			KeyPair         _reciKP   = CMSTestUtil.makeKeyPair();
			X509Certificate _reciCert = CMSTestUtil.makeCertificate(_reciKP, _reciDN, _signKP, _signDN);
			
			CMSEnvelopedDataGenerator edGen = new CMSEnvelopedDataGenerator();

			edGen.addKeyTransRecipient(_reciCert);

			CMSEnvelopedData ed = edGen.generate(
                                    new CMSProcessableByteArray(data),
                                    CMSEnvelopedDataGenerator.AES192_CBC, "BC");

            RecipientInformationStore  recipients = ed.getRecipientInfos();

            Collection  c = recipients.getRecipients();
            Iterator    it = c.iterator();

            while (it.hasNext())
            {
                RecipientInformation   recipient = (RecipientInformation)it.next();

                byte[] recData = recipient.getContent(_reciKP.getPrivate(), "BC");

                assertEquals(true, Arrays.equals(data, recData));
            }
		}
		catch(Exception ex) {
			log(ex);
			fail();
		}
	}

	public void testKeyTransAES256()
    {
		try {
			byte[]          data     = "WallaWallaWashington".getBytes();

			String          _signDN   = "O=Bouncy Castle, C=AU";
			KeyPair         _signKP   = CMSTestUtil.makeKeyPair();  
			X509Certificate _signCert = CMSTestUtil.makeCertificate(_signKP, _signDN, _signKP, _signDN);

			String          _origDN   = "CN=Bob, OU=Sales, O=Bouncy Castle, C=AU";
			KeyPair         _origKP   = CMSTestUtil.makeKeyPair();
			X509Certificate _origCert = CMSTestUtil.makeCertificate(_origKP, _origDN, _signKP, _signDN);

			String          _reciDN   = "CN=Doug, OU=Sales, O=Bouncy Castle, C=AU";
			KeyPair         _reciKP   = CMSTestUtil.makeKeyPair();
			X509Certificate _reciCert = CMSTestUtil.makeCertificate(_reciKP, _reciDN, _signKP, _signDN);
			
			CMSEnvelopedDataGenerator edGen = new CMSEnvelopedDataGenerator();

			edGen.addKeyTransRecipient(_reciCert);

			CMSEnvelopedData ed = edGen.generate(
                                    new CMSProcessableByteArray(data),
                                    CMSEnvelopedDataGenerator.AES256_CBC, "BC");

            RecipientInformationStore  recipients = ed.getRecipientInfos();

            Collection  c = recipients.getRecipients();
            Iterator    it = c.iterator();

            while (it.hasNext())
            {
                RecipientInformation   recipient = (RecipientInformation)it.next();

                byte[] recData = recipient.getContent(_reciKP.getPrivate(), "BC");

                assertEquals(true, Arrays.equals(data, recData));
            }
		}
		catch(Exception ex) {
			log(ex);
			fail();
		}
	}

	public void testKeyTransRC4()
    {
		try
        {
			byte[]          data     = "WallaWallaBouncyCastle".getBytes();

			String          _signDN   = "O=Bouncy Castle, C=AU";
			KeyPair         _signKP   = CMSTestUtil.makeKeyPair();  
			X509Certificate _signCert = CMSTestUtil.makeCertificate(_signKP, _signDN, _signKP, _signDN);

			String          _origDN   = "CN=Bob, OU=Sales, O=Bouncy Castle, C=AU";
			KeyPair         _origKP   = CMSTestUtil.makeKeyPair();
			X509Certificate _origCert = CMSTestUtil.makeCertificate(_origKP, _origDN, _signKP, _signDN);

			String          _reciDN   = "CN=Doug, OU=Sales, O=Bouncy Castle, C=AU";
			KeyPair         _reciKP   = CMSTestUtil.makeKeyPair();
			X509Certificate _reciCert = CMSTestUtil.makeCertificate(_reciKP, _reciDN, _signKP, _signDN);
			
			CMSEnvelopedDataGenerator edGen = new CMSEnvelopedDataGenerator();

			edGen.addKeyTransRecipient(_reciCert);

			CMSEnvelopedData ed = edGen.generate(
                                    new CMSProcessableByteArray(data),
                                    "1.2.840.113549.3.4", "BC");

            RecipientInformationStore  recipients = ed.getRecipientInfos();

            Collection  c = recipients.getRecipients();
            Iterator    it = c.iterator();

            while (it.hasNext())
            {
                RecipientInformation   recipient = (RecipientInformation)it.next();

                byte[] recData = recipient.getContent(_reciKP.getPrivate(), "BC");

                assertEquals(true, Arrays.equals(data, recData));
            }
		}
		catch(Exception ex)
        {
			log(ex);
			fail();
		}
	}

	public void testKeyTransODES()
    {
		try
        {
			byte[]          data     = "WallaWallaBouncyCastle".getBytes();

			String          _signDN   = "O=Bouncy Castle, C=AU";
			KeyPair         _signKP   = CMSTestUtil.makeKeyPair();  
			X509Certificate _signCert = CMSTestUtil.makeCertificate(_signKP, _signDN, _signKP, _signDN);

			String          _origDN   = "CN=Bob, OU=Sales, O=Bouncy Castle, C=AU";
			KeyPair         _origKP   = CMSTestUtil.makeKeyPair();
			X509Certificate _origCert = CMSTestUtil.makeCertificate(_origKP, _origDN, _signKP, _signDN);

			String          _reciDN   = "CN=Doug, OU=Sales, O=Bouncy Castle, C=AU";
			KeyPair         _reciKP   = CMSTestUtil.makeKeyPair();
			X509Certificate _reciCert = CMSTestUtil.makeCertificate(_reciKP, _reciDN, _signKP, _signDN);
			
			CMSEnvelopedDataGenerator edGen = new CMSEnvelopedDataGenerator();

			edGen.addKeyTransRecipient(_reciCert);

			CMSEnvelopedData ed = edGen.generate(
                                    new CMSProcessableByteArray(data),
                                    "1.3.14.3.2.7", "BC");

            RecipientInformationStore  recipients = ed.getRecipientInfos();

            Collection  c = recipients.getRecipients();
            Iterator    it = c.iterator();

            while (it.hasNext())
            {
                RecipientInformation   recipient = (RecipientInformation)it.next();

                byte[] recData = recipient.getContent(_reciKP.getPrivate(), "BC");

                assertEquals(true, Arrays.equals(data, recData));
            }
		}
		catch(Exception ex)
        {
			log(ex);
			fail();
        }
    }

	public void testDESKEK()
    {
		try {
			byte[]    data = "WallaWallaWashington".getBytes();
			SecretKey kek  = CMSTestUtil.makeDesede192Key();
			
			CMSEnvelopedDataGenerator edGen = new CMSEnvelopedDataGenerator();

            byte[]  kekId = new byte[] { 1, 2, 3, 4, 5 };

			edGen.addKEKRecipient(kek, kekId);

			CMSEnvelopedData ed = edGen.generate(
                                    new CMSProcessableByteArray(data),
                                    CMSEnvelopedDataGenerator.DES_EDE3_CBC, "BC");

            RecipientInformationStore  recipients = ed.getRecipientInfos();

            Collection  c = recipients.getRecipients();
            Iterator    it = c.iterator();

            while (it.hasNext())
            {
                RecipientInformation   recipient = (RecipientInformation)it.next();

                byte[] recData = recipient.getContent(kek, "BC");

                assertEquals(true, Arrays.equals(data, recData));
            }
			
		}
		catch(Exception ex) {
			log(ex);
			fail();
		}
	}

	public void testRC2KEK()
    {
		try {
			byte[]    data = "WallaWallaWashington".getBytes();
			SecretKey kek  = CMSTestUtil.makeRC2128Key();
			
			CMSEnvelopedDataGenerator edGen = new CMSEnvelopedDataGenerator();

            byte[]  kekId = new byte[] { 1, 2, 3, 4, 5 };

			edGen.addKEKRecipient(kek, kekId);

			CMSEnvelopedData ed = edGen.generate(
                                    new CMSProcessableByteArray(data),
                                    CMSEnvelopedDataGenerator.DES_EDE3_CBC, "BC");

            RecipientInformationStore  recipients = ed.getRecipientInfos();

            Collection  c = recipients.getRecipients();
            Iterator    it = c.iterator();

            while (it.hasNext())
            {
                RecipientInformation   recipient = (RecipientInformation)it.next();

                byte[] recData = recipient.getContent(kek, "BC");

                assertEquals(true, Arrays.equals(data, recData));
            }
			
		}
		catch(Exception ex) {
			log(ex);
			fail();
		}
	}
}
