/*
 * Syncany, www.syncany.org
 * Copyright (C) 2011-2016 Philipp C. Heckel <philipp.heckel@gmail.com>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package org.syncany.crypto;

import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.spec.InvalidKeySpecException;
import java.util.Arrays;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.logging.Level;
import java.util.logging.Logger;

import javax.crypto.SecretKey;

import org.syncany.util.StringUtil;

/**
 * The cipher session is used by the {@link MultiCipherOutputStream} and the
 * {@link MultiCipherInputStream} to reference the application's master key,
 * and to temporarily store and retrieve derived secret keys.
 *
 * <p>While a cipher session does not create a master key, it creates and manages
 * the derived keys using the {@link CipherUtil} class:
 *
 * <ul>
 *   <li>Keys used by {@link MultiCipherOutputStream} (for writing new files) are
 *       reused a number of times before a new salted key is created. The main
 *       purpose of reusing keys is to increase performance. Because the master
 *       key is cryptographically strong, the derived keys can be reused a few
 *       times without any drawbacks on security. The class keeps one secret key
 *       per {@link CipherSpec}.
 *
 *   <li>Keys used by {@link MultiCipherInputStream} (when reading files) are
 *       cached in order to minimize the amount of keys that have to be created when
 *       files are processed.
 * </ul>
 *
 * @author Philipp C. Heckel <philipp.heckel@gmail.com>
 */
public class CipherSession {
	private static final Logger logger = Logger.getLogger(CipherSession.class.getSimpleName());
	private static final int DEFAULT_SECRET_KEY_READ_CACHE_SIZE = 20;
	private static final int DEFAULT_SECRET_KEY_WRITE_REUSE_COUNT = 100;

	private SecretKey masterKey;

	private Map<CipherSpecWithSalt, SecretKeyCacheEntry> secretKeyReadCache;
	private int secretKeyReadCacheSize;

	private Map<CipherSpec, SecretKeyCacheEntry> secretKeyWriteCache;
	private int secretKeyWriteReuseCount;

	/**
	 * Creates a new cipher session, using the given master key. Derived keys will be created
	 * from that master key.
	 *
	 * <p>The default settings for read/write key cache will be used. Refer to
	 * {@link CipherSession the class description} for more details. Default values:
	 * {@link #DEFAULT_SECRET_KEY_READ_CACHE_SIZE} and {@link #DEFAULT_SECRET_KEY_WRITE_REUSE_COUNT}
	 *
	 * @param masterKey The master key, used for deriving new read/write keys
	 */
	public CipherSession(SaltedSecretKey masterKey) {
		this(masterKey, DEFAULT_SECRET_KEY_READ_CACHE_SIZE, DEFAULT_SECRET_KEY_WRITE_REUSE_COUNT);
	}

	/**
	 * Creates a new cipher session, using the given master key. Derived keys will be created
	 * from that master key.
	 *
	 * <p>This method expects a reuse-count for write keys and a cache size for the read-key cache.
	 * Refer to {@link CipherSession the class description} for more details.
	 *
	 * @param masterKey The master key, used for deriving new read/write
	 * @param secretKeyReadCacheSize Number of read keys to store in the cache (higher means more performance, but more memory usage)
	 * @param secretKeyWriteReuseCount Number of times to reuse a write key (higher means more performance, but lower security)
	 */
	public CipherSession(SaltedSecretKey masterKey, int secretKeyReadCacheSize, int secretKeyWriteReuseCount) {
		this.masterKey = masterKey;

		this.secretKeyReadCache = new LinkedHashMap<CipherSpecWithSalt, SecretKeyCacheEntry>();
		this.secretKeyReadCacheSize = secretKeyReadCacheSize;

		this.secretKeyWriteCache = new HashMap<CipherSpec, SecretKeyCacheEntry>();
		this.secretKeyWriteReuseCount = secretKeyWriteReuseCount;
	}

	/**
	 * Returns the master key
	 */
	public SecretKey getMasterKey() {
		return masterKey;
	}

	/**
	 * Creates or retrieves a derived write secret key and updates the cache and re-use count. If a key is reused more
	 * than the threshold defined in {@link #secretKeyWriteReuseCount} (as set in {@link #CipherSession(SaltedSecretKey, int, int) the constructor},
	 * a new key is created and added to the cache.
	 *
	 * <p>If a new key needs to be created, {@link CipherUtil} is used to do so.
	 *
	 * <p>Contrary to the read cache, the write cache key is a only {@link CipherSpec}, i.e. only one secret key
	 * per cipher spec can be held in the cache.
	 *
	 * @param cipherSpec Defines the type of key to be created (or retrieved); used as key for the cache retrieval
	 * @return Returns a newly created secret key or a cached key
	 * @throws Exception If an error occurs with key creation
	 */
	public SaltedSecretKey getWriteSecretKey(CipherSpec cipherSpec) throws Exception {
		SecretKeyCacheEntry secretKeyCacheEntry = secretKeyWriteCache.get(cipherSpec);

		// Remove key if use more than X times
		if (secretKeyCacheEntry != null && secretKeyCacheEntry.getUseCount() >= secretKeyWriteReuseCount) {
			logger.log(Level.FINE, "- Removed WRITE secret key from cache, because it was used " + secretKeyCacheEntry.getUseCount() + " times.");

			secretKeyWriteCache.remove(cipherSpec);
			secretKeyCacheEntry = null;
		}

		// Return cached key, or create a new one
		if (secretKeyCacheEntry != null) {
			secretKeyCacheEntry.increaseUseCount();

			logger.log(Level.FINE, "- Using CACHED WRITE secret key " + secretKeyCacheEntry.getSaltedSecretKey().getAlgorithm() + ", with salt "
					+ StringUtil.toHex(secretKeyCacheEntry.getSaltedSecretKey().getSalt()));
			return secretKeyCacheEntry.getSaltedSecretKey();
		}
		else {
			SaltedSecretKey saltedSecretKey = createSaltedSecretKey(cipherSpec);

			secretKeyCacheEntry = new SecretKeyCacheEntry(saltedSecretKey);
			secretKeyWriteCache.put(cipherSpec, secretKeyCacheEntry);

			logger.log(Level.FINE, "- Created NEW WRITE secret key " + secretKeyCacheEntry.getSaltedSecretKey().getAlgorithm()
					+ ", and added to cache, with salt " + StringUtil.toHex(saltedSecretKey.getSalt()));
			return saltedSecretKey;
		}
	}

	/**
	 * Creates a new secret key or retrieves it from the read cache. If the given cipher spec / salt combination
	 * is found in the cache, the cached secret key is returned. If not, a new key is created. Keys are removed
	 * from the cache when the cache reached the size defined by {@link #secretKeyReadCacheSize} (as set in
	 * {@link #CipherSession(SaltedSecretKey, int, int) the constructor}.
	 *
	 * <p>If a new key needs to be created, {@link CipherUtil} is used to do so.
	 *
	 * <p>Contrary to the write cache, the read cache key is a combination of {@link CipherSpec} and a salt. For
	 * each cipher spec, multiple salted keys can reside in the cache at the same time.
	 *
	 * @param cipherSpec Defines the type of key to be created (or retrieved); used as one part of the key for cache retrieval
	 * @param salt Defines the salt for the key to be created (or retrieved); used as one part of the key for cache retrieval
	 * @return Returns a newly created secret key or a cached key
	 * @throws Exception If an error occurs with key creation
	 */
	public SaltedSecretKey getReadSecretKey(CipherSpec cipherSpec, byte[] salt) throws Exception {
		CipherSpecWithSalt cipherSpecWithSalt = new CipherSpecWithSalt(cipherSpec, salt);
		SecretKeyCacheEntry secretKeyCacheEntry = secretKeyReadCache.get(cipherSpecWithSalt);

		if (secretKeyCacheEntry != null) {
			logger.log(Level.FINE, "- Using CACHED READ secret key " + secretKeyCacheEntry.getSaltedSecretKey().getAlgorithm() + ", with salt "
					+ StringUtil.toHex(salt));
			return secretKeyCacheEntry.getSaltedSecretKey();
		}
		else {
			if (secretKeyReadCache.size() >= secretKeyReadCacheSize) {
				CipherSpecWithSalt firstKey = secretKeyReadCache.keySet().iterator().next();
				secretKeyReadCache.remove(firstKey);

				logger.log(Level.FINE, "- Removed oldest READ secret key from cache.");
			}

			SaltedSecretKey saltedSecretKey = createSaltedSecretKey(cipherSpec, salt);
			secretKeyCacheEntry = new SecretKeyCacheEntry(saltedSecretKey);

			secretKeyReadCache.put(cipherSpecWithSalt, secretKeyCacheEntry);

			logger.log(Level.FINE, "- Created NEW READ secret key " + secretKeyCacheEntry.getSaltedSecretKey().getAlgorithm()
					+ ", and added to cache, with salt " + StringUtil.toHex(salt));
			return saltedSecretKey;
		}
	}

	private SaltedSecretKey createSaltedSecretKey(CipherSpec cipherSpec) throws InvalidKeySpecException, NoSuchAlgorithmException,
	NoSuchProviderException {
		byte[] salt = CipherUtil.createRandomArray(MultiCipherOutputStream.SALT_SIZE);
		return createSaltedSecretKey(cipherSpec, salt);
	}

	private SaltedSecretKey createSaltedSecretKey(CipherSpec cipherSpec, byte[] salt) throws InvalidKeySpecException, NoSuchAlgorithmException,
	NoSuchProviderException {
		return CipherUtil.createDerivedKey(masterKey, salt, cipherSpec);
	}

	private static class SecretKeyCacheEntry {
		private SaltedSecretKey saltedSecretKey;
		private int useCount;

		public SecretKeyCacheEntry(SaltedSecretKey saltedSecretKey) {
			this.saltedSecretKey = saltedSecretKey;
			this.useCount = 1;
		}

		public SaltedSecretKey getSaltedSecretKey() {
			return saltedSecretKey;
		}

		public int getUseCount() {
			return useCount;
		}

		public void increaseUseCount() {
			useCount++;
		}
	}

	private static class CipherSpecWithSalt {
		private CipherSpec cipherSpec;
		private byte[] salt;

		public CipherSpecWithSalt(CipherSpec cipherSpec, byte[] salt) {
			this.cipherSpec = cipherSpec;
			this.salt = salt;
		}

		@Override
		public int hashCode() {
			final int prime = 31;
			int result = 1;
			result = prime * result + ((cipherSpec == null) ? 0 : cipherSpec.hashCode());
			result = prime * result + Arrays.hashCode(salt);
			return result;
		}

		@Override
		public boolean equals(Object obj) {
			if (this == obj) {
				return true;
			}
			if (obj == null) {
				return false;
			}
			if (!(obj instanceof CipherSpecWithSalt)) {
				return false;
			}
			CipherSpecWithSalt other = (CipherSpecWithSalt) obj;
			if (cipherSpec == null) {
				if (other.cipherSpec != null) {
					return false;
				}
			}
			else if (!cipherSpec.equals(other.cipherSpec)) {
				return false;
			}
			if (!Arrays.equals(salt, other.salt)) {
				return false;
			}
			return true;
		}
	}

}
