/*
 * sshIt: simple, powerful, open-source SSH client for Android
 * Copyright 2014 Alpha LLC
 *
 * 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
 * limitations under the License.
 */

package org.alphallc.sshit.util;

/**
 * This class is from:
 *
 * Encryptor.java
 * Copyright 2008 Zach Scrivena
 * zachscrivena@gmail.com
 * http://zs.freeshell.org/
 */

import java.security.MessageDigest;
import java.security.SecureRandom;
import java.util.Arrays;

import javax.crypto.Cipher;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;


/**
 * Perform AES-128 encryption.
 */
public final class Encryptor
{
	/** name of the character set to use for converting between characters and bytes */
	private static final String CHARSET_NAME = "UTF-8";

	/** random number generator algorithm */
	private static final String RNG_ALGORITHM = "SHA1PRNG";

	/** message digest algorithm (must be sufficiently long to provide the key and initialization vector) */
	private static final String DIGEST_ALGORITHM = "SHA-256";

	/** key algorithm (must be compatible with CIPHER_ALGORITHM) */
	private static final String KEY_ALGORITHM = "AES";

	/** cipher algorithm (must be compatible with KEY_ALGORITHM) */
	private static final String CIPHER_ALGORITHM = "AES/CBC/PKCS5Padding";


	/**
	* Private constructor that should never be called.
	*/
	private Encryptor()
	{}


	/**
	* Encrypt the specified cleartext using the given password.
	* With the correct salt, number of iterations, and password, the decrypt() method reverses
	* the effect of this method.
	* This method generates and uses a random salt, and the user-specified number of iterations
	* and password to create a 16-byte secret key and 16-byte initialization vector.
	* The secret key and initialization vector are then used in the AES-128 cipher to encrypt
	* the given cleartext.
	*
	* @param salt
	*	  salt that was used in the encryption (to be populated)
	* @param iterations
	*	  number of iterations to use in salting
	* @param password
	*	  password to be used for encryption
	* @param cleartext
	*	  cleartext to be encrypted
	* @return
	*	  ciphertext
	* @throws Exception
	*	  on any error encountered in encryption
	*/
	public static byte[] encrypt(
			final byte[] salt,
			final int iterations,
			final String password,
			final byte[] cleartext)
			throws Exception
	{
		/* generate salt randomly */
		SecureRandom.getInstance(RNG_ALGORITHM).nextBytes(salt);

		/* compute key and initialization vector */
		final MessageDigest shaDigest = MessageDigest.getInstance(DIGEST_ALGORITHM);
		byte[] pw = password.getBytes(CHARSET_NAME);

		for (int i = 0; i < iterations; i++)
		{
			/* add salt */
			final byte[] salted = new byte[pw.length + salt.length];
			System.arraycopy(pw, 0, salted, 0, pw.length);
			System.arraycopy(salt, 0, salted, pw.length, salt.length);
			Arrays.fill(pw, (byte) 0x00);

			/* compute SHA-256 digest */
			shaDigest.reset();
			pw = shaDigest.digest(salted);
			Arrays.fill(salted, (byte) 0x00);
		}

		/* extract the 16-byte key and initialization vector from the SHA-256 digest */
		final byte[] key = new byte[16];
		final byte[] iv = new byte[16];
		System.arraycopy(pw, 0, key, 0, 16);
		System.arraycopy(pw, 16, iv, 0, 16);
		Arrays.fill(pw, (byte) 0x00);

		/* perform AES-128 encryption */
		final Cipher cipher = Cipher.getInstance(CIPHER_ALGORITHM);

		cipher.init(
				Cipher.ENCRYPT_MODE,
				new SecretKeySpec(key, KEY_ALGORITHM),
				new IvParameterSpec(iv));

		Arrays.fill(key, (byte) 0x00);
		Arrays.fill(iv, (byte) 0x00);

		return cipher.doFinal(cleartext);
	}


	/**
	* Decrypt the specified ciphertext using the given password.
	* With the correct salt, number of iterations, and password, this method reverses the effect
	* of the encrypt() method.
	* This method uses the user-specified salt, number of iterations, and password
	* to recreate the 16-byte secret key and 16-byte initialization vector.
	* The secret key and initialization vector are then used in the AES-128 cipher to decrypt
	* the given ciphertext.
	*
	* @param salt
	*	  salt to be used in decryption
	* @param iterations
	*	  number of iterations to use in salting
	* @param password
	*	  password to be used for decryption
	* @param ciphertext
	*	  ciphertext to be decrypted
	* @return
	*	  cleartext
	* @throws Exception
	*	  on any error encountered in decryption
	*/
	public static byte[] decrypt(
			final byte[] salt,
			final int iterations,
			final String password,
			final byte[] ciphertext)
			throws Exception
	{
		/* compute key and initialization vector */
		final MessageDigest shaDigest = MessageDigest.getInstance(DIGEST_ALGORITHM);
		byte[] pw = password.getBytes(CHARSET_NAME);

		for (int i = 0; i < iterations; i++)
		{
			/* add salt */
			final byte[] salted = new byte[pw.length + salt.length];
			System.arraycopy(pw, 0, salted, 0, pw.length);
			System.arraycopy(salt, 0, salted, pw.length, salt.length);
			Arrays.fill(pw, (byte) 0x00);

			/* compute SHA-256 digest */
			shaDigest.reset();
			pw = shaDigest.digest(salted);
			Arrays.fill(salted, (byte) 0x00);
		}

		/* extract the 16-byte key and initialization vector from the SHA-256 digest */
		final byte[] key = new byte[16];
		final byte[] iv = new byte[16];
		System.arraycopy(pw, 0, key, 0, 16);
		System.arraycopy(pw, 16, iv, 0, 16);
		Arrays.fill(pw, (byte) 0x00);

		/* perform AES-128 decryption */
		final Cipher cipher = Cipher.getInstance(CIPHER_ALGORITHM);

		cipher.init(
				Cipher.DECRYPT_MODE,
				new SecretKeySpec(key, KEY_ALGORITHM),
				new IvParameterSpec(iv));

		Arrays.fill(key, (byte) 0x00);
		Arrays.fill(iv, (byte) 0x00);

		return cipher.doFinal(ciphertext);
	}
}
