/* $Id$ */
package org.noreply.fancydress.crypto;

import java.security.SecureRandom;
import org.bouncycastle.crypto.digests.SHA1Digest;
import org.bouncycastle.crypto.engines.AESEngine;
import org.bouncycastle.crypto.params.KeyParameter;
import org.noreply.fancydress.misc.Util;

/**
 * This class implements all cryptographic primitives that are needed for a
 * type III implementation.
 */
public class CryptoPrimitives {
	/**
	 * length (in octets) of a symmetric AES-128 key.
	 */
	public static final int KEY_LEN = 16;
	/**
	 * length (in octets) of our SPRP (LIONESS) key.
	 */
	public static final int SPRP_KEY_LEN = 20;
	/**
	 * length (in octets) of our hash function's (SHA-1) output.
	 */
	public static final int HASH_LEN = 20;

	/**
	 * A secure pseudo random number generator.
	 */
	public static SecureRandom secureRandom = new SecureRandom();
	/**
	 * Our SHA-1 Digest.
	 */
	public static SHA1Digest digest = new SHA1Digest();
	/**
	 * The AES Engine used for symmetric operations.
	 */
	public static AESEngine aes_cipher = new AESEngine();

	/**
	 * XOR two octet arrays.
	 *
	 * @param a octet array 1
	 * @param b octet array 2
	 * @return an octet array of equal length where x[i] = a[i] XOR b[i] for all i
	 * @throws IllegalArgumentException if a and b do not have the same length
	 */
	public static byte[] xor(byte[] a, byte[] b) {
		if (a.length != b.length)
			throw new IllegalArgumentException("Arguments to xor must have the same length");
		byte[] result = new byte[a.length];
		for (int i=0; i<a.length; i++)
			result[i] = (byte) (a[i] ^ b[i]);
		return(result);
	}
	/**
	 * Get <code>n</code> random octets from our secure pseudo random number generator.
	 *
	 * @param n number of octets
	 * @return an octet array of n octets of your fines randomness
	 */
	public static byte[] rand(int n) {
		byte[] result = new byte[n];
		secureRandom.nextBytes(result);
		return(result);
	}
	/**
	 * Get <code>n</code> octets of zeroes.
	 *
	 * @param n number of octets
	 * @return an octet array of n octets of zeroes
	 */
	public static byte[] zero(int n) {
		byte[] result = new byte[n];
		return(result);
	}
	/**
	 * Hash the message provided in <code>m</code>.
	 *
	 * @param m a message
	 * @return the SHA-1 hash of <code>m</code>
	 */
	public static byte[] hash(byte[] m) {
		digest.update(m, 0, m.length);
		byte[] result = new byte[HASH_LEN];
		digest.doFinal(result, 0);
		return(result);
	}
	/**
	 * Hash the result of the concatenation of <code>m1</code> and <code>m2</code>.
	 *
	 * @param m1 message 1
	 * @param m2 message 2
	 * @return the SHA-1 hash of <code>m1 + m2</code>
	 */
	public static byte[] hash(byte[] m1, byte[] m2) {
		digest.update(m1, 0, m1.length);
		digest.update(m2, 0, m2.length);
		byte[] result = new byte[HASH_LEN];
		digest.doFinal(result, 0);
		return(result);
	}
	/**
	 * Hash the result of the concatenation of <code>m1</code>, <code>m2</code> and <code>m3</code>.
	 *
	 * @param m1 message 1
	 * @param m2 message 2
	 * @param m3 message 3
	 * @return the SHA-1 hash of <code>m1 + m2 + m3</code>
	 */
	public static byte[] hash(byte[] m1, byte[] m2, byte[] m3) {
		digest.update(m1, 0, m1.length);
		digest.update(m2, 0, m2.length);
		digest.update(m3, 0, m3.length);
		byte[] result = new byte[HASH_LEN];
		digest.doFinal(result, 0);
		return(result);
	}
	/**
	 * Return <code>n</code> "random" octets generated based on key <code>k</code>.
	 *
	 * Generates <code>n</code> "random" octets or a key stream based on
	 * <code>k</code>. The stream is generated by using AES in counter mode
	 * with key <code>k</code>.
	 * 
	 * @param k the key for AES counter mode
	 * @param n number of octets requested
	 * @return a keystream of n octets
	 * @throws IllegalArgumentException if <code>k</code> is not KEY_LEN (=16) octets long.
	 */
	public static byte[] prng(byte[] k, int n) {
		byte[] result = new byte[ n ];
		byte[] block = new byte[ KEY_LEN ];
		byte[] encrypted = new byte[ KEY_LEN ];
		if (k.length != KEY_LEN)
			throw new IllegalArgumentException("Argument k to encrypt must be KEY_LEN bytes");
		KeyParameter params = new KeyParameter(k);
		aes_cipher.init(true, params);
		for (int i=0, p=0; i<n; p++) {
			block[block.length - 1] = (byte) (p & 0xff);
			block[block.length - 2] = (byte) ((p >> 8) & 0xff);
			block[block.length - 3] = (byte) ((p >> 16) & 0xff);
			block[block.length - 4] = (byte) ((p >> 24) & 0xff);

			aes_cipher.processBlock(block, 0, encrypted, 0);
			for (int j=0; j<encrypted.length && i<n; j++, i++)
				result[i] = encrypted[j];
		};
		return(result);
	}
	/**
	 * Encrypt a message <code>m</code> using AES counter mode with key <code>k</code>.
	 *
	 * Returns the result of xor(m, prng(k, len(m))), i.e. the message <code>m</code>
	 * encrypted using AES in counter mode with key <code>k</code>.
	 *
	 * @param k the key for AES counter mode
	 * @param m the message to encrypt
	 * @return the encrypted message. (len(e) == len(m))
	 * @throws IllegalArgumentException if <code>k</code> is not KEY_LEN (=16) octets long.
	 */
	public static byte[] encrypt(byte[] k, byte[] m) {
		byte[] result;
		if (k.length != KEY_LEN)
			throw new IllegalArgumentException("Argument k to encrypt must be KEY_LEN bytes");
		result = xor(m, prng(k, m.length));
		return(result);
	}
	/**
	 * Encrypt a message using our super pseudorandom permutation.
	 *
	 * Encrypt message <code>m</code> using an instance of the LIONESS
	 * super pseudorandom permutation (SPRP).
	 *
	 * This SPRP has the property that any change in the encrypted value
	 * will make the decryption look like random bits
	 *
	 * @param k the key for our SPRP function
	 * @param m the message to encrypt
	 * @return the encrypted message. (len(e) == len(m))
	 * @throws IllegalArgumentException if <code>k</code> is not SPRP_KEY_LEN (=20) octets long
	 * @throws IllegalArgumentException if <code>m</code> is shorter than SPRP_KEY_LEN (=20) octets
	 */
	public static byte[] sprpEncrypt(byte[] k, byte[] m) {
		if (m.length < HASH_LEN)
			throw new IllegalArgumentException("Argument m to sprp_encrypt must be of at least HASH_LEN bytes");
		if (k.length != SPRP_KEY_LEN)
			throw new IllegalArgumentException("Argument k to sprp_encrypt must be SPRP_KEY_LEN bytes");
		byte[] result;
		byte[] k1 = new byte[HASH_LEN];
		byte[] k2 = new byte[HASH_LEN];
		byte[] k3 = new byte[HASH_LEN];
		byte[] k4 = new byte[HASH_LEN];

		System.arraycopy(k, 0, k1, 0, k.length);
		System.arraycopy(k, 0, k2, 0, k.length);
		System.arraycopy(k, 0, k3, 0, k.length);
		System.arraycopy(k, 0, k4, 0, k.length);
		k2[19] = (byte) (k1[19] ^ 0x01);
		k3[19] = (byte) (k1[19] ^ 0x02);
		k4[19] = (byte) (k1[19] ^ 0x03);
		byte[] l = Util.slice(m, 0, HASH_LEN);
		byte[] r = Util.slice(m, HASH_LEN, m.length-HASH_LEN);
		r = encrypt( Util.slice(hash( k1, l, k1), 0, KEY_LEN), r);
		l = xor( l, hash(k2, r, k2 ));
		r = encrypt( Util.slice(hash( k3, l, k3), 0, KEY_LEN), r);
		l = xor( l, hash(k4, r, k4 ));
		result = Util.concat(l, r);
		return(result);
	}
	/**
	 * Encrypt a message using our super pseudorandom permutation.
	 *
	 * The message will be encrypted using a subkey constructed from
	 * <code>k</code> and <code>p</code> (=HASH(k | p)).
	 *
	 * @param k a master key for the SPRP encryption.
	 * @param p a string to build a subkey with.
	 * @param m the message to encrypt
	 * @return the encrypted message. (len(e) == len(m))
	 * @throws IllegalArgumentException if <code>m</code> is shorter than SPRP_KEY_LEN (=20) octets
	 * @see #sprpEncrypt(byte[] k, byte[] m)
	 */
	public static byte[] sprpEncrypt(byte[] k, String p, byte[] m) {
		return sprpEncrypt(hash(k, Util.toOctets(p)) ,m);
	}
	/**
	 * Decrypt a message using our super pseudorandom permutation.
	 *
	 * Decrypt message <code>m</code> using our instance of the LIONESS
	 * super pseudorandom permutation (SPRP).
	 *
	 * @param k the key for our SPRP function
	 * @param m the message to decrypt
	 * @return the encrypted message. (len(e) == len(m))
	 * @throws IllegalArgumentException if <code>k</code> is not SPRP_KEY_LEN (=20) octets long
	 * @throws IllegalArgumentException if <code>m</code> is shorter than SPRP_KEY_LEN (=20) octets
	 * @see #sprpEncrypt(byte[] k, byte[] m)
	 */
	public static byte[] sprpDecrypt(byte[] k, byte[] m) {
		if (m.length < HASH_LEN)
			throw new IllegalArgumentException("Argument m to sprp_decrypt must be of at least HASH_LEN bytes");
		if (k.length != SPRP_KEY_LEN)
			throw new IllegalArgumentException("Argument k to sprp_decrypt must be SPRP_KEY_LEN bytes");
		byte[] result;
		byte[] k1 = new byte[HASH_LEN];
		byte[] k2 = new byte[HASH_LEN];
		byte[] k3 = new byte[HASH_LEN];
		byte[] k4 = new byte[HASH_LEN];

		System.arraycopy(k, 0, k1, 0, k.length);
		System.arraycopy(k, 0, k2, 0, k.length);
		System.arraycopy(k, 0, k3, 0, k.length);
		System.arraycopy(k, 0, k4, 0, k.length);
		k2[19] = (byte) (k1[19] ^ 0x01);
		k3[19] = (byte) (k1[19] ^ 0x02);
		k4[19] = (byte) (k1[19] ^ 0x03);
		byte[] l = Util.slice(m, 0, HASH_LEN);
		byte[] r = Util.slice(m, HASH_LEN, m.length-HASH_LEN);
		l = xor( l, hash(k4, r, k4 ));
		r = encrypt( Util.slice(hash( k3, l, k3), 0, KEY_LEN), r);
		l = xor( l, hash(k2, r, k2 ));
		r = encrypt( Util.slice(hash( k1, l, k1), 0, KEY_LEN), r);
		result = Util.concat(l, r);
		return(result);
	}
	/**
	 * Decrypt a message using our super pseudorandom permutation.
	 *
	 * The message will be decrypted using a subkey constructed from
	 * <code>k</code> and <code>p</code> (=HASH(k | p)).
	 *
	 * @param k a master key for the SPRP decryption.
	 * @param p a string to build a subkey with.
	 * @param m the message to decrypt
	 * @return the encrypted message. (len(e) == len(m))
	 * @throws IllegalArgumentException if <code>m</code> is shorter than SPRP_KEY_LEN (=20) octets
	 * @see #sprpEncrypt(byte[] k, byte[] m)
	 * @see #sprpDecrypt(byte[] k, byte[] m)
	 */
	public static byte[] sprpDecrypt(byte[] k, String p, byte[] m) {
		return sprpDecrypt(hash(k, Util.toOctets(p)) ,m);
	}
	/**
	 * Build a subkey from <code>k</code> and <code>p</code>.
	 *
	 * The subkey is generated using HASH(k | p).
	 *
	 * @param k a master key
	 * @param p a string to build a subkey with
	 * @return the subkey
	 * @throws IllegalArgumentException if <code>k</code> is not KEY_LEN (=16) octets long
	 */
	public static byte[] subKey(byte[] k, String p) {
		if (k.length != KEY_LEN)
			throw new IllegalArgumentException("Argument k to subKey must be KEY_LEN bytes");
		return Util.slice(hash(k, Util.toOctets(p)), 0, 16);
	}
}