package org.bouncycastle.crypto.kems;

import java.math.BigInteger;
import java.security.SecureRandom;

import org.bouncycastle.crypto.CipherParameters;
import org.bouncycastle.crypto.DerivationFunction;
import org.bouncycastle.crypto.KeyEncapsulation;
import org.bouncycastle.crypto.params.KDFParameters;
import org.bouncycastle.crypto.params.KeyParameter;
import org.bouncycastle.crypto.params.RSAKeyParameters;
import org.bouncycastle.util.BigIntegers;

/**
 * The RSA Key Encapsulation Mechanism (RSA-KEM) from ISO 18033-2.
 */
public class RSAKeyEncapsulation
    implements KeyEncapsulation
{
    private static final BigInteger ZERO = BigInteger.valueOf(0);
    private static final BigInteger ONE = BigInteger.valueOf(1);

    private DerivationFunction kdf;
    private SecureRandom rnd;
    private RSAKeyParameters key;

    /**
     * Set up the RSA-KEM.
     *
     * @param kdf the key derivation function to be used.
     * @param rnd the random source for the session key.
     */
    public RSAKeyEncapsulation(
        DerivationFunction kdf,
        SecureRandom rnd)
    {
        this.kdf = kdf;
        this.rnd = rnd;
    }


    /**
     * Initialise the RSA-KEM.
     *
     * @param key the recipient's public (for encryption) or private (for decryption) key.
     */
    public void init(CipherParameters key)
        throws IllegalArgumentException
    {
        if (!(key instanceof RSAKeyParameters))
        {
            throw new IllegalArgumentException("RSA key required");
        }
        else
        {
            this.key = (RSAKeyParameters)key;
        }
    }


    /**
     * Generate and encapsulate a random session key.
     *
     * @param out    the output buffer for the encapsulated key.
     * @param outOff the offset for the output buffer.
     * @param keyLen the length of the random session key.
     * @return the random session key.
     */
    public CipherParameters encrypt(byte[] out, int outOff, int keyLen)
        throws IllegalArgumentException
    {
        if (key.isPrivate())
        {
            throw new IllegalArgumentException("Public key required for encryption");
        }

        BigInteger n = key.getModulus();
        BigInteger e = key.getExponent();

        // Generate the ephemeral random and encode it    
        BigInteger r = BigIntegers.createRandomInRange(ZERO, n.subtract(ONE), rnd);
        byte[] R = BigIntegers.asUnsignedByteArray((n.bitLength() + 7) / 8, r);

        // Encrypt the random and encode it     
        BigInteger c = r.modPow(e, n);
        byte[] C = BigIntegers.asUnsignedByteArray((n.bitLength() + 7) / 8, c);
        System.arraycopy(C, 0, out, outOff, C.length);


        // Initialise the KDF
        kdf.init(new KDFParameters(R, null));

        // Generate the secret key
        byte[] K = new byte[keyLen];
        kdf.generateBytes(K, 0, K.length);

        return new KeyParameter(K);
    }


    /**
     * Generate and encapsulate a random session key.
     *
     * @param out    the output buffer for the encapsulated key.
     * @param keyLen the length of the random session key.
     * @return the random session key.
     */
    public CipherParameters encrypt(byte[] out, int keyLen)
    {
        return encrypt(out, 0, keyLen);
    }


    /**
     * Decrypt an encapsulated session key.
     *
     * @param in     the input buffer for the encapsulated key.
     * @param inOff  the offset for the input buffer.
     * @param inLen  the length of the encapsulated key.
     * @param keyLen the length of the session key.
     * @return the session key.
     */
    public CipherParameters decrypt(byte[] in, int inOff, int inLen, int keyLen)
        throws IllegalArgumentException
    {
        if (!key.isPrivate())
        {
            throw new IllegalArgumentException("Private key required for decryption");
        }

        BigInteger n = key.getModulus();
        BigInteger d = key.getExponent();

        // Decode the input
        byte[] C = new byte[inLen];
        System.arraycopy(in, inOff, C, 0, C.length);
        BigInteger c = new BigInteger(1, C);

        // Decrypt the ephemeral random and encode it
        BigInteger r = c.modPow(d, n);
        byte[] R = BigIntegers.asUnsignedByteArray((n.bitLength() + 7) / 8, r);

        // Initialise the KDF
        kdf.init(new KDFParameters(R, null));

        // Generate the secret key
        byte[] K = new byte[keyLen];
        kdf.generateBytes(K, 0, K.length);

        return new KeyParameter(K);
    }

    /**
     * Decrypt an encapsulated session key.
     *
     * @param in     the input buffer for the encapsulated key.
     * @param keyLen the length of the session key.
     * @return the session key.
     */
    public CipherParameters decrypt(byte[] in, int keyLen)
    {
        return decrypt(in, 0, in.length, keyLen);
    }
}

@@ -1133,33 +1133,32 @@
 d = rnd;%0A    %7D%0A%0A
-%0A
     /**%0A     * I
@@ -1515,35 +1515,9 @@
   %7D%0A
-        else%0A        %7B%0A    
+%0A
     
@@ -1562,20 +1562,9 @@
     
-    %7D%0A    %7D%0A
+%7D
 %0A%0A  
@@ -2340,88 +2340,8 @@
 nd);
-%0A        byte%5B%5D R = BigIntegers.asUnsignedByteArray((n.bitLength() + 7) / 8, r);
 %0A%0A  
@@ -2562,25 +2562,24 @@
 ngth);%0A%0A
-%0A
         
 // Initi
@@ -2574,226 +2574,39 @@
     
-// Initialise the KDF%0A        kdf.init(new KDFParameters(R, null));%0A%0A        // Generate the secret key%0A        byte%5B%5D K = new byte%5BkeyLen%5D;%0A        kdf.generateBytes(K, 0, K.length);%0A%0A        return new KeyParameter(K
+return generateKey(n, r, keyLen
 );%0A 
@@ -2603,33 +2603,32 @@
 keyLen);%0A    %7D%0A%0A
-%0A
     /**%0A     * G
@@ -2964,17 +2964,16 @@
 %0A    %7D%0A%0A
-%0A
     /**%0A
@@ -3891,32 +3891,33 @@
 dPow(d, n);%0A
+%0A
         
 byte%5B%5D R = B
@@ -3908,307 +3908,39 @@
     
-byte%5B%5D R = BigIntegers.asUnsignedByteArray((n.bitLength() + 7) / 8, r);%0A%0A        // Initialise the KDF%0A        kdf.init(new KDFParameters(R, null));%0A%0A        // Generate the secret key%0A        byte%5B%5D K = new byte%5BkeyLen%5D;%0A        kdf.generateBytes(K, 0, K.length);%0A%0A        return new KeyParameter(K
+return generateKey(n, r, keyLen
 );%0A 
@@ -4280,11 +4280,413 @@
 );%0A    %7D
+%0A%0A    protected KeyParameter generateKey(BigInteger n, BigInteger r, int keyLen)%0A    %7B%0A        byte%5B%5D R = BigIntegers.asUnsignedByteArray((n.bitLength() + 7) / 8, r);%0A%0A        // Initialise the KDF%0A        kdf.init(new KDFParameters(R, null));%0A%0A        // Generate the secret key%0A        byte%5B%5D K = new byte%5BkeyLen%5D;%0A        kdf.generateBytes(K, 0, K.length);%0A%0A        return new KeyParameter(K);%0A    %7D
 %0A%7D%0A

Factor out method generateKey