package org.postgresql.ssl;

import java.io.File;
import java.io.FileInputStream;
import java.io.InputStream;
import java.security.KeyStore;
import java.util.Arrays;
import java.util.Properties;

import javax.net.ssl.KeyManager;
import javax.net.ssl.KeyManagerFactory;
import javax.net.ssl.SSLContext;
import javax.net.ssl.TrustManager;
import javax.net.ssl.TrustManagerFactory;

public class ValidatingFactory extends WrappedFactory {

	/**
	 * Template string for the property name which indicates the name of a
	 * key/trust store file.
	 */
	private static final String SSL__FILE = "ssl._.file";

	/**
	 * Template string for the property name which indicates the type of
	 * key/trust store to instanciate.
	 */
	private static final String SSL__TYPE = "ssl._.type";

	/**
	 * Template String for the property name which indicates a password to a
	 * key/trust store.
	 */
	private static final String SSL__PASSWORD = "ssl._.password";

	/**
	 * Template String for the property name which indicates the algorithm to
	 * use when instanciateing a Key/TrustManagerFactory. This is somewhat
	 * redundant since the default algorithm may be set via your security
	 * provider.
	 */
	private static final String SSL__ALGORITHM = "ssl._.algorithm";

	/**
	 * String used to differenciate a keystore from a truststore in the Property
	 * templates.
	 */
	private static final String KEYSTORE = "keystore";

	/**
	 * Property used to indicate the keystore file. If this property does
	 * not exist in the Properties object passed to the consturctor, then the
	 * file System.getProperty("user.home")/postgresql/.postgresql.jks will be
	 * used instead.
	 */
	public static final String SSL_KEYSTORE_FILE = SSL__FILE.replaceFirst("_",
			KEYSTORE);

	/**
	 * Property used to indicate the type of keystore which will be loaded. The
	 * supported values depend on your security provider. KeyStore.getDefaultType()
	 * is used, if this property is not present.
	 *
	 * @see KeyStore#getDefaultType()
	 */
	public static final String SSL_KEYSTORE_TYPE = SSL__TYPE.replaceFirst("_",
			KEYSTORE);

	/**
	 * Property used to indicate the password which will be used to load the key
	 * store. If this option is not present, then an empty password will be
	 * used. Note that some keystore tools will not allow you to create a keystore
	 * without a password.
	 */
	public static final String SSL_KEYSTORE_PASSWORD = SSL__PASSWORD
			.replaceFirst("_", KEYSTORE);

	/**
	 * Property used to indicate the algorithm which should be used when
	 * creating the KeyManagerFactory, which will in turn supply the KeyManagers
	 * to the SSLContext. If this property is not present, then
	 * KeyManagerFactory.getDefaultAlgorithm will be used instead.
	 * 
	 * @see KeyManagerFactory#getDefaultAlgorithm()
	 */
	public static final String SSL_KEYSTORE_ALGORITHM = SSL__ALGORITHM
			.replaceFirst("_", KEYSTORE);

	/**
	 * The property name used to indicate that the default JRE key managers
	 * should be used. The existance of this property in the Properties object
	 * passed to the constuctor indicates that the default should be used, and
	 * the value is ignored.
	 * <p>
	 * This option will have the affect of passing a null argument to the
	 * SSLContext.init method. This behavior is useful if you are connecting to
	 * a server which does not require client authentication.
	 * <p>
	 * All other SSL_KEYSTORE_* options will be ignored if this option is set.
	 * 
	 * @see SSLContext#init(KeyManager[], TrustManager[],
	 *      java.security.SecureRandom)
	 */
	public static final String SSL_USE_DEFAULT_KEY_MANAGER = "ssl.use.default.key.manager";

	/**
	 * String used to differenciate a truststore from a keystore in the Property
	 * templates.
	 */
	private static final String TRUSTSTORE = "truststore";

	/**
	 * Property used to indicate the struststore file. If this property does
	 * not exist in the Properties object passed to the consturctor, then the
	 * file System.getProperty("user.home")/postgresql/.postgresql.jks will be
	 * used instead.
	 */
	public static final String SSL_TRUSTSTORE_FILE = SSL__FILE.replaceFirst(
			"_", TRUSTSTORE);

	/**
	 * Property used to indicate the type of keystore which will be loaded. The
	 * supported values depend on your security provider. KeyStore.getDefaultTYpe
	 * is used, if this property is not present.
	 *
	 * @see KeyStore#getDefaultType()
	 */
	public static final String SSL_TRUSTSTORE_TYPE = SSL__TYPE.replaceFirst(
			"_", TRUSTSTORE);

	/**
	 * Property used to indicate the password which will be used to load the key
	 * store. If this option is not present, then an empty password will be
	 * used. A trust store commonly does not contain secure information, so 
	 * it is likely that this option is not required.
	 */
	public static final String SSL_TRUSTSTORE_PASSWORD = SSL__PASSWORD
			.replaceFirst("_", TRUSTSTORE);

	/**
	 * Property used to indicate the algorithm which should be used when
	 * creating the TrustManagerFactory, which will in turn supply the TrustManagers
	 * to the SSLContext. If this property is not present, then
	 * TrustManagerFactory.getDefaultAlgorithm will be used instead.
	 * 
	 * @see TrustManagerFactory#getDefaultAlgorithm()
	 */
	public static final String SSL_TRUSTSTORE_ALGORITHM = SSL__ALGORITHM
			.replaceFirst("_", TRUSTSTORE);

	/**
	 * The property name used to indicate that the default JRE trust managers
	 * should be used. The existance of this property in the Properties object
	 * passed to the constuctor indicates that the default should be used, and
	 * the value is ignored.
	 * <p>
	 * This option will have the affect of passing a null argument to the
	 * SSLContext.init method. This behavior is useful if you are connecting to
	 * a server which has a certificate signed by a widely accepted Certificate
	 * Authority. See the java API for details.
	 * <p>
	 * All other SSL_TRUSTSTORE_* options will be ignored if this option is set.
	 * 
	 * @see SSLContext#init(KeyManager[], TrustManager[],
	 *      java.security.SecureRandom)
	 */
	public static final String SSL_USE_DEFAULT_TRUST_MANAGER = "ssl.use.default.trust.manager";

	/**
	 * A Property used to indicate that when constructing the trust store, the
	 * information provided for the key store should be used. This existance of
	 * the property is used to indicate that this behavior is desired, and the
	 * value is ignored.
	 * <p>
	 * This is useful in cases where the the root server certificate is stored
	 * in the client's keystore, which is a common case.
	 */
	public static final String SSL_KEYSTORE_IS_TRUSTSTORE = "ssl.keystore.is.truststore";

	/**
	 * A Property name used to identify the ssl protocol to use to connect to 
	 * the server. The supported values depends on your security provider. If
	 * no option is specified, then SSLv3 is used.
	 */
	public static final String SSL_PROTOCOL = "ssl.protocol";

	/**
	 * A Property used to indicate a custom security provider.
	 */
	public static final String SSL_PROVIDER = "ssl.provider";

	/**
	 * Constructor which uses the property values specified by the
	 * public static values of this class to retrieve ssl factory
	 * configuration options.
	 * 
	 * @param props
	 * @throws Exception This constructor can fail for a wide 
	 *                   variety of reasons, all of which result in some
	 *                   sort of Exception.
	 */
	public ValidatingFactory(Properties props) throws Exception {

		// Obtain the trust managers for the ssl context.
		TrustManager[] trustManagers = null;
		if (props.getProperty(SSL_USE_DEFAULT_TRUST_MANAGER) == null) {
			trustManagers = createTrustManagers(props);
		}

		// Obtain the key managers for the ssl context.
		KeyManager[] keyManagers = null;
		if (props.getProperty(SSL_USE_DEFAULT_KEY_MANAGER) == null) {
			keyManagers = createKeyManagers(props);
		}

		// Create the ssl context.
		String protocol = getProtocol(props);
		SSLContext sslContext = null;
		if (props.getProperty(SSL_PROVIDER) != null) {
			sslContext = SSLContext.getInstance(protocol, props
					.getProperty(SSL_PROVIDER));
		} else {
			sslContext = SSLContext.getInstance(protocol);
		}

		// initialize the context.
		sslContext.init(keyManagers, trustManagers, null);

		// Create the wrapped socket factory.
		_factory = sslContext.getSocketFactory();
	}

	/**
	 * @param props The properties object passed to the constructor should
	 *              be passed to thie method.
	 * @return Array of TrustManagers which should be used to initialize
	 * 				the SSLContext.
	 * @throws Exception
	 */
	protected static TrustManager[] createTrustManagers(Properties props)
			throws Exception {
		// We may want to simply used the key store for trusted certificates.
		String mode = TRUSTSTORE;
		if (props.getProperty(SSL_KEYSTORE_IS_TRUSTSTORE) != null) {
			mode = KEYSTORE;
		}

		// Obtain the keystore.
		KeyStore trustStore = getLoadedKeyStore(props, mode);

		
		String provider = props.getProperty(SSL_PROVIDER);
		String algorithm = getAlgorithm(props, mode);
		TrustManagerFactory tmf = null;
		
		if (provider != null) {
			tmf = TrustManagerFactory.getInstance(algorithm, provider);
		} else {
			tmf = TrustManagerFactory.getInstance(algorithm);
		}

		tmf.init(trustStore);

		return tmf.getTrustManagers();
	}

	/**
	 * @param props Argument should be the same as the one passed to
	 *              the constructor.
	 * @return A KeyManager array intended for use for the SSLContext.init method.
	 * @throws Exception Many reasons this could fail...
	 */
	private static KeyManager[] createKeyManagers(Properties props)
			throws Exception {
		KeyStore trustStore = getLoadedKeyStore(props, KEYSTORE);

		String provider = props.getProperty(SSL_PROVIDER);
		String algorithm = getAlgorithm(props, KEYSTORE);
		KeyManagerFactory kmf = null;
		if (provider != null) {
			kmf = KeyManagerFactory.getInstance(algorithm, provider);
		} else {
			kmf = KeyManagerFactory.getInstance(algorithm);
		}

		char[] password = getPWD(props, KEYSTORE);
		try {
			kmf.init(trustStore, password);
		}finally {
			Arrays.fill(password, '\0');
		}
		
		return kmf.getKeyManagers();
	}

	private static String getAlgorithm(Properties props, String mode) {
		String ans = props.getProperty(SSL__ALGORITHM.replaceFirst("_", mode));
		if (ans != null) {
			return ans;
		}

		if (mode.equals(KEYSTORE)) {
			return KeyManagerFactory.getDefaultAlgorithm();
		} else {
			return TrustManagerFactory.getDefaultAlgorithm();
		}
	}

	private static KeyStore getLoadedKeyStore(Properties props, String mode)
			throws Exception {
		File keystoreFile = getStoreFile(props, mode);

		String keyStoreType = getStoreType(props, mode, keystoreFile);

		String provider = props.getProperty(SSL_PROVIDER);
		KeyStore store = null;
		if (provider != null) {
			store = KeyStore.getInstance(keyStoreType, provider);
		} else {
			store = KeyStore.getInstance(keyStoreType);
		}

		InputStream strm = new FileInputStream(keystoreFile);
		char[] password = getPWD(props, mode);
		try {
			store.load(strm, password);
		} finally {
			Arrays.fill(password, '\0');			
		}

		strm.close();
		
		return store;
	}

	private static File getStoreFile(Properties props, String mode) {
		String ks = props.getProperty(SSL__FILE.replaceFirst("_", mode));
		if (ks != null) {
			return new File(ks);
		}

		File ans = new File(System.getProperty("user.home"));
		ans = new File(ans, ".postgresql");
		ans = new File(ans, "postgresql.jks");
		return ans;
	}

	private static String getStoreType(Properties props, String mode, File keyStore) {
		String kst = props.getProperty(SSL__TYPE.replaceFirst("_", mode));
		if (kst != null) {
			return kst;
		}

		return KeyStore.getDefaultType();
	}

	private static char[] getPWD(Properties props, String mode) {
		// Unfortunately, the password is in a string. We don't really
		// have anyway of wiping it.
		String pwd = props.getProperty(SSL__PASSWORD.replaceFirst("_", mode),
				"");
		return pwd.toCharArray();
	}

	private static String getProtocol(Properties props) {
		if (props.containsKey(SSL_PROTOCOL)) {
			return props.getProperty(SSL_PROTOCOL);
		}
		return "SSLv3";
	}
}