/*
 * Copyright (c) 2012-2013 Christian Plattner. All rights reserved.
 * Please refer to the LICENSE.txt for licensing details.
 */

package ch.ethz.ssh2;

import java.io.CharArrayWriter;
import java.io.File;
import java.io.FileReader;
import java.io.IOException;
import java.net.Socket;
import java.util.List;
import java.util.ArrayList;

import ch.ethz.ssh2.crypto.CryptoWishList;
import ch.ethz.ssh2.crypto.PEMDecoder;
import ch.ethz.ssh2.server.ServerConnectionState;
import java.security.KeyPair;
import java.security.PrivateKey;
import java.security.interfaces.DSAPrivateKey;
import java.security.interfaces.ECPrivateKey;
import java.security.interfaces.RSAPrivateKey;
import ch.ethz.ssh2.transport.ServerTransportManager;

/**
 * A server-side SSH-2 connection.
 *
 * @author Christian
 *
 */
public class ServerConnection
{
	/**
	 * The softwareversion presented to the SSH-2 client.
	 */
	private String softwareversion = String.format("Ganymed_SSHD_%s", Version.getSpecification());

	private final ServerConnectionState state = new ServerConnectionState(this);

	/**
	 * Creates a new <code>ServerConnection</code> that will communicate
	 * with the client over the given <code>Socket</code>.
	 * <p>
	 * Note: you need to call {@link #connect()} or {@link #connect(int)} to
	 * perform the initial handshake and establish the encrypted communication.
	 *
	 * @see #connect(int)
	 *
	 * @param s The socket
	 */
	public ServerConnection(Socket s)
	{
		this(s, null, null, null);
	}

    public ServerConnection(Socket s, String softwareversion) {
        this(s, null, null, null);
        this.softwareversion = softwareversion;
    }

    /**
	 * Creates a new <code>ServerConnection</code> that will communicate
	 * with the client over the given <code>Socket</code>.
	 * <p>
	 * Note: you need to call {@link #connect()} or {@link #connect(int)} to
	 * perform the initial handshake and establish the encrypted communication.
	 * <p>
	 * Please read the javadoc for the {@link #connect(int)} method.
	 *
	 * @see #connect(int)
	 *
	 * @param s The socket
	 * @param dsa_key The DSA hostkey, may be <code>NULL</code>
	 * @param rsa_key The RSA hostkey, may be <code>NULL</code>
	 * @param ec_key  The EC  hostkey, may be <code>NULL</code>
	 */
	public ServerConnection(Socket s, KeyPair dsa_key, KeyPair rsa_key, KeyPair ec_key)
	{
		state.s = s;
		state.softwareversion = softwareversion;
		state.next_dsa_key = dsa_key;
		state.next_rsa_key = rsa_key;
		state.next_ec_key  = ec_key;
		fixCryptoWishList(state.next_cryptoWishList, state.next_dsa_key, state.next_rsa_key, state.next_ec_key);
	}

	/**
	 * Establish the connection and block until the first handshake has completed.
	 * <p>
	 * Note: this is a wrapper that calls <code>connect(0)</code> (i.e., connect with no timeout).
	 * <p>
	 * Please read the javadoc for the {@link #connect(int)} method.
	 *
	 * @see #connect(int)
	 *
	 * @throws IOException
	 */
	public synchronized void connect() throws IOException
	{
		connect(0);
	}

	/**
	 * Establish the connection and block until the first handshake has completed.
	 * <p>
	 * Note 1: at least one DSA, RSA or EC hostkey must be set before calling this method.
	 * <p>
	 * Note 2: You must set the callbacks for authentication ({@link #setAuthenticationCallback(ServerAuthenticationCallback)})
	 * and connection events ({@link #setServerConnectionCallback(ServerConnectionCallback)}).
	 *
	 * @see #setPEMHostKey(char[], String)
	 * @see #setPEMHostKey(File, String)
	 * @see #setRsaHostKey(RSAPrivateKey)
	 * @see #setDsaHostKey(DSAPrivateKey)
	 *
	 * @param timeout_milliseconds Timeout in milliseconds, <code>0</code> means no timeout.
	 * @throws IOException
	 */
	public synchronized void connect(int timeout_milliseconds) throws IOException
	{
		synchronized (state)
		{
			if (state.cb_conn == null)
				throw new IllegalStateException("The callback for connection events has not been set.");

			if (state.cb_auth == null)
				throw new IllegalStateException("The callback for authentication events has not been set.");

			if (state.tm != null)
				throw new IllegalStateException("The initial handshake has already been started.");

			if ((state.next_dsa_key == null) && (state.next_rsa_key == null) && (state.next_ec_key == null))
				throw new IllegalStateException("Neither an RSA nor a DSA nor an EC host key has been specified!");

			state.tm = new ServerTransportManager(state.s);
		}

		state.tm.connect(state);

		/* Wait until first KEX has finished */

		state.tm.getConnectionInfo(1);
	}

	/**
	 * Retrieve the underlying socket.
	 *
	 * @return the socket that has been passed to the constructor.
	 */
	public Socket getSocket()
	{
		return state.s;
	}

	/**
	 * Force an asynchronous key re-exchange (the call does not block). The
	 * latest values set for MAC, Cipher and DH group exchange parameters will
	 * be used. If a key exchange is currently in progress, then this method has
	 * the only effect that the so far specified parameters will be used for the
	 * next (client driven) key exchange. You may call this method only after
	 * the initial key exchange has been established.
	 * <p>
	 * Note: This implementation will never start automatically a key exchange (other than the initial one)
	 * unless you or the connected SSH-2 client ask for it.
	 *
	 * @throws IOException
	 *             In case of any failure behind the scenes.
	 */
	public synchronized void forceKeyExchange() throws IOException
	{
		synchronized (state)
		{
			if (state.tm == null)
				throw new IllegalStateException(
						"Cannot force another key exchange, you need to start the key exchange first.");

			state.tm.forceKeyExchange(state.next_cryptoWishList, null, state.next_dsa_key, state.next_rsa_key, state.next_ec_key);
		}
	}

	/**
	 * Returns a {@link ConnectionInfo} object containing the details of
	 * the connection. May be called as soon as the first key exchange has been
	 * started. The method blocks in case the first key exchange has not been completed.
	 * <p>
	 * Note: upon return of this method, authentication may still be pending.
	 *
	 * @return A {@link ConnectionInfo} object.
	 * @throws IOException
	 *             In case of any failure behind the scenes; e.g., first key exchange was aborted.
	 */
	public synchronized ConnectionInfo getConnectionInfo() throws IOException
	{
		synchronized (state)
		{
			if (state.tm == null)
				throw new IllegalStateException(
						"Cannot get details of connection, you need to start the key exchange first.");
		}

		return state.tm.getConnectionInfo(1);
	}

	/**
	 * Change the current DSA hostkey. Either a DSA or RSA or EC private key must be set for a successful handshake with
	 * the client.
	 * <p>
	 * Note: You can change an existing DSA hostkey after the initial kex exchange (the new value will
	 * be used during the next server initiated key exchange), but you cannot remove (i.e., set to <code>null</code>) the
	 * current DSA key, otherwise the next key exchange may fail in case the client supports only DSA hostkeys.
	 *
	 * @param dsa_hostkey
	 */
	public synchronized void setDsaHostKey(KeyPair dsa_hostkey)
	{
		synchronized (state)
		{
			if ((dsa_hostkey == null) && (state.next_dsa_key != null) && (state.tm != null))
				throw new IllegalStateException("Cannot remove DSA hostkey after first key exchange.");

			state.next_dsa_key = dsa_hostkey;
			fixCryptoWishList(state.next_cryptoWishList, state.next_dsa_key, state.next_rsa_key, state.next_ec_key);
		}
	}

	/**
	 * Change the current RSA hostkey. Either a DSA or RSA or EC private key must be set for a successful handshake with
	 * the client.
	 * <p>
	 * Note: You can change an existing RSA hostkey after the initial kex exchange (the new value will
	 * be used during the next server initiated key exchange), but you cannot remove (i.e., set to <code>null</code>) the
	 * current RSA key, otherwise the next key exchange may fail in case the client supports only RSA hostkeys.
	 *
	 * @param rsa_hostkey
	 */
	public synchronized void setRsaHostKey(KeyPair rsa_hostkey)
	{
		synchronized (state)
		{
			if ((rsa_hostkey == null) && (state.next_rsa_key != null) && (state.tm != null))
				throw new IllegalStateException("Cannot remove RSA hostkey after first key exchange.");

			state.next_rsa_key = rsa_hostkey;
			fixCryptoWishList(state.next_cryptoWishList, state.next_dsa_key, state.next_rsa_key, state.next_ec_key);
		}
	}

	/**
	 * Change the current EC hostkey. Either a DSA or RSA or EC private key must be set for a successful handshake with
	 * the client.
	 * <p>
	 * Note: You can change an existing EC hostkey after the initial kex exchange (the new value will
	 * be used during the next server initiated key exchange), but you cannot remove (i.e., set to <code>null</code>) the
	 * current EC key, otherwise the next key exchange may fail in case the client supports only EC hostkeys.
	 *
	 * @param rsa_hostkey
	 */
	public synchronized void setEcHostKey(KeyPair ec_hostkey)
	{
		synchronized (state)
		{
			if ((ec_hostkey == null) && (state.next_ec_key != null) && (state.tm != null))
				throw new IllegalStateException("Cannot remove EC hostkey after first key exchange.");

			state.next_ec_key = ec_hostkey;
			fixCryptoWishList(state.next_cryptoWishList, state.next_dsa_key, state.next_rsa_key, state.next_ec_key);
		}
	}

	/**
	 * Utility method that loads a PEM based hostkey (either RSA or DSA based) and
	 * calls either <code>setRsaHostKey()</code> or <code>setDsaHostKey()</code>.
	 *
	 * @param pemdata The PEM data
	 * @param password Password, may be null in case the PEM data is not password protected
	 * @throws IOException In case of any error.
	 */
	public void setPEMHostKey(char[] pemdata, String password) throws IOException
	{
		KeyPair pair = PEMDecoder.decode(pemdata, password);
        PrivateKey key = pair.getPrivate();

		if (key instanceof DSAPrivateKey) setDsaHostKey(pair);

		if (key instanceof RSAPrivateKey) setRsaHostKey(pair);

        if (key instanceof ECPrivateKey) setEcHostKey(pair);
	}

	/**
	 * Utility method that loads a hostkey from a PEM file (either RSA or DSA based) and
	 * calls either <code>setRsaHostKey()</code> or <code>setDsaHostKey()</code>.
	 *
	 * @param pemFile The PEM file
	 * @param password Password, may be null in case the PEM file is not password protected
	 * @throws IOException
	 */
	public void setPEMHostKey(File pemFile, String password) throws IOException
	{
		if (pemFile == null)
			throw new IllegalArgumentException("pemfile argument is null");

		char[] buff = new char[256];

		CharArrayWriter cw = new CharArrayWriter();

		FileReader fr = new FileReader(pemFile);

		while (true)
		{
			int len = fr.read(buff);
			if (len < 0)
				break;
			cw.write(buff, 0, len);
		}

		fr.close();

		setPEMHostKey(cw.toCharArray(), password);
	}

	private void fixCryptoWishList(CryptoWishList next_cryptoWishList, KeyPair next_dsa_key, KeyPair next_rsa_key, KeyPair next_ec_key)
	{
        List<String> algos = new ArrayList<String>();
		if (next_ec_key != null)  algos.add("ecdsa-sha2-nistp521");
		if (next_ec_key != null)  algos.add("ecdsa-sha2-nistp384");
		if (next_ec_key != null)  algos.add("ecdsa-sha2-nistp256");
		if (next_dsa_key != null) algos.add("ssh-dss");
		if (next_rsa_key != null) algos.add("ssh-rsa");
 	    next_cryptoWishList.serverHostKeyAlgorithms = new String[algos.size()];
        algos.toArray(next_cryptoWishList.serverHostKeyAlgorithms);
	}

	/**
	 * Callback interface with methods that will be called upon events
	 * generated by the client (e.g., client opens a new Session which results in a <code>ServerSession</code>).
	 * <p>
	 * Note: This must be set before the first handshake.
	 *
	 * @param cb The callback implementation
	 */
	public synchronized void setServerConnectionCallback(ServerConnectionCallback cb)
	{
		synchronized (state)
		{
			state.cb_conn = cb;
		}
	}

	/**
	 * Callback interface with methods that will be called upon authentication events.
	 * <p>
	 * Note: This must be set before the first handshake.
	 *
	 * @param cb The callback implementation
	 */
	public synchronized void setAuthenticationCallback(ServerAuthenticationCallback cb)
	{
		synchronized (state)
		{
			state.cb_auth = cb;
		}
	}

	/**
	 * Close the connection to the SSH-2 server. All assigned sessions will be
	 * closed, too. Can be called at any time. Don't forget to call this once
	 * you don't need a connection anymore - otherwise the receiver thread may
	 * run forever.
	 */
	public void close()
	{
        synchronized (state)
      		{
      			if (state.cm != null)
      				state.cm.closeAllChannels();

      			if (state.tm != null)
      			{
      				state.tm.close();
      			}
      		}
	}

	public void close(IOException t)
	{
		synchronized (state)
		{
			if (state.cm != null)
				state.cm.closeAllChannels();

			if (state.tm != null)
			{
				state.tm.close(t);
			}
		}
	}
}