/*
 * Copyright (c) 2006-2011 Christian Plattner. All rights reserved.
 * Please refer to the LICENSE.txt for licensing details.
 */
package ch.ethz.ssh2.transport;

import java.io.IOException;
import java.io.InterruptedIOException;
import java.security.DigestException;
import java.security.KeyPair;
import java.security.PrivateKey;
import java.security.SecureRandom;
import java.security.interfaces.DSAPrivateKey;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.ECPrivateKey;
import java.util.Arrays;
import java.util.Set;
import java.util.TreeSet;

import ch.ethz.ssh2.ConnectionInfo;
import ch.ethz.ssh2.DHGexParameters;
import ch.ethz.ssh2.compression.CompressionFactory;
import ch.ethz.ssh2.compression.Compressor;
import ch.ethz.ssh2.crypto.CryptoWishList;
import ch.ethz.ssh2.crypto.KeyMaterial;
import ch.ethz.ssh2.crypto.cipher.BlockCipher;
import ch.ethz.ssh2.crypto.cipher.BlockCipherFactory;
import ch.ethz.ssh2.crypto.digest.MAC;
import ch.ethz.ssh2.log.Logger;
import ch.ethz.ssh2.packets.PacketKexInit;
import ch.ethz.ssh2.packets.PacketNewKeys;
import java.security.interfaces.DSAPrivateKey;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.ECPrivateKey;

/**
 * @version $Id: KexManager.java 152 2014-04-28 11:02:23Z dkocher@sudo.ch $
 */
public abstract class KexManager implements MessageHandler {
    protected static final Logger log = Logger.getLogger(KexManager.class);

    private static final Set<String> HOSTKEY_ALGS = new TreeSet<String>();
    static {
        HOSTKEY_ALGS.add("ecdsa-sha2-nistp256");
        HOSTKEY_ALGS.add("ecdsa-sha2-nistp384");
        HOSTKEY_ALGS.add("ecdsa-sha2-nistp521");
        HOSTKEY_ALGS.add("ssh-rsa");
        HOSTKEY_ALGS.add("ssh-dss");
    }

    private static final Set<String> KEX_ALGS = new TreeSet<String>();
    static {
        KEX_ALGS.add("ecdh-sha2-nistp256");
        KEX_ALGS.add("ecdh-sha2-nistp384");
        KEX_ALGS.add("ecdh-sha2-nistp521");
        KEX_ALGS.add("diffie-hellman-group-exchange-sha256");
        KEX_ALGS.add("diffie-hellman-group-exchange-sha1");
        KEX_ALGS.add("diffie-hellman-group14-sha1");
        KEX_ALGS.add("diffie-hellman-group1-sha1");
    }

    KexState kxs;
    int kexCount = 0;
    KeyMaterial km;
    byte[] sessionId;
    ClientServerHello csh;

    final Object accessLock = new Object();
    ConnectionInfo lastConnInfo = null;

    boolean connectionClosed = false;

    boolean ignore_next_kex_packet = false;

    final TransportManager tm;

    CryptoWishList nextKEXcryptoWishList;
    DHGexParameters nextKEXdhgexParameters;
    KeyPair nextKEXdsakey;
    KeyPair nextKEXrsakey;
    KeyPair nextKEXeckey;

    final SecureRandom rnd;

    public KexManager(TransportManager tm, ClientServerHello csh, CryptoWishList initialCwl, SecureRandom rnd) {
        this.tm = tm;
        this.csh = csh;
        this.nextKEXcryptoWishList = initialCwl;
        this.nextKEXdhgexParameters = new DHGexParameters();
        this.rnd = rnd;
    }

    public ConnectionInfo getOrWaitForConnectionInfo(int minKexCount) throws IOException {
        synchronized (accessLock) {
            while (true) {
                if ((lastConnInfo != null) && (lastConnInfo.keyExchangeCounter >= minKexCount)) {
                    return lastConnInfo;
                }

                if (connectionClosed) {
                    throw(IOException) new IOException("Key exchange was not finished, connection is closed.").initCause(tm.getReasonClosedCause());
                }

                try {
                    accessLock.wait();
                }
                catch (InterruptedException e) {
                    throw new InterruptedIOException(e.getMessage());
                }
            }
        }
    }

    private String getFirstMatch(String[] client, String[] server) throws NegotiateException {
        if (client == null || server == null) {
            throw new IllegalArgumentException();
        }

        for (String c : client) {
            for (String s : server) {
                if (c.equals(s)) {
                    return c;
                }
            }
        }

        throw new NegotiateException(String.format("Negotiation failed for %s", Arrays.toString(server)));
    }

    private boolean compareFirstOfNameList(String[] a, String[] b) {
        if (a == null || b == null) {
            throw new IllegalArgumentException();
        }

        if ((a.length == 0) && (b.length == 0)) {
            return true;
        }

        if ((a.length == 0) || (b.length == 0)) {
            return false;
        }

        return (a[0].equals(b[0]));
    }

    private boolean isGuessOK(KexParameters cpar, KexParameters spar) {
        if (cpar == null || spar == null) {
            throw new IllegalArgumentException();
        }

        if (!compareFirstOfNameList(cpar.kex_algorithms, spar.kex_algorithms)) {
            return false;
        }

        if (!compareFirstOfNameList(cpar.server_host_key_algorithms, spar.server_host_key_algorithms)) {
            return false;
        }

        /*
         * We do NOT check here if the other algorithms can be agreed on, this
         * is just a check if kex_algorithms and server_host_key_algorithms were
         * guessed right!
         */
        return true;
    }

    protected NegotiatedParameters mergeKexParameters(KexParameters client, KexParameters server)
    throws NegotiateException {
        NegotiatedParameters np = new NegotiatedParameters();
        np.kex_algo = getFirstMatch(client.kex_algorithms, server.kex_algorithms);
        log.info("kex_algo=" + np.kex_algo);
        np.server_host_key_algo = getFirstMatch(client.server_host_key_algorithms,
                                                server.server_host_key_algorithms);
        log.info("server_host_key_algo=" + np.server_host_key_algo);
        np.enc_algo_client_to_server = getFirstMatch(client.encryption_algorithms_client_to_server,
                                       server.encryption_algorithms_client_to_server);
        np.enc_algo_server_to_client = getFirstMatch(client.encryption_algorithms_server_to_client,
                                       server.encryption_algorithms_server_to_client);
        log.info("enc_algo_client_to_server=" + np.enc_algo_client_to_server);
        log.info("enc_algo_server_to_client=" + np.enc_algo_server_to_client);
        np.mac_algo_client_to_server = getFirstMatch(client.mac_algorithms_client_to_server,
                                       server.mac_algorithms_client_to_server);
        np.mac_algo_server_to_client = getFirstMatch(client.mac_algorithms_server_to_client,
                                       server.mac_algorithms_server_to_client);
        log.info("mac_algo_client_to_server=" + np.mac_algo_client_to_server);
        log.info("mac_algo_server_to_client=" + np.mac_algo_server_to_client);
        np.comp_algo_client_to_server = getFirstMatch(client.compression_algorithms_client_to_server,
                                        server.compression_algorithms_client_to_server);
        np.comp_algo_server_to_client = getFirstMatch(client.compression_algorithms_server_to_client,
                                        server.compression_algorithms_server_to_client);
        log.info("comp_algo_client_to_server=" + np.comp_algo_client_to_server);
        log.info("comp_algo_server_to_client=" + np.comp_algo_server_to_client);
        np.lang_client_to_server = getFirstMatch(client.languages_client_to_server,
                                   server.languages_client_to_server);
        np.lang_server_to_client = getFirstMatch(client.languages_server_to_client,
                                   server.languages_server_to_client);

        if (isGuessOK(client, server)) {
            np.guessOK = true;
        }

        return np;
    }

    public synchronized void initiateKEX(CryptoWishList cwl, DHGexParameters dhgex, KeyPair dsa, KeyPair rsa, KeyPair ec)
    throws IOException {
        nextKEXcryptoWishList = cwl;
        nextKEXdhgexParameters = dhgex;
        nextKEXdsakey = dsa;
        nextKEXrsakey = rsa;
        nextKEXeckey  = ec;

        if (kxs == null) {
            kxs = new KexState();
            kxs.local_dsa_key = dsa;
            kxs.local_rsa_key = rsa;
            kxs.local_ec_key  = ec;
            kxs.dhgexParameters = nextKEXdhgexParameters;
            kxs.localKEX = new PacketKexInit(nextKEXcryptoWishList, rnd);
            tm.sendKexMessage(kxs.localKEX.getPayload());
        }
    }

    private boolean establishKeyMaterial() throws IOException {
        try {
            int mac_cs_key_len = MAC.getKeyLen(kxs.np.mac_algo_client_to_server);
            int enc_cs_key_len = BlockCipherFactory.getKeySize(kxs.np.enc_algo_client_to_server);
            int enc_cs_block_len = BlockCipherFactory.getBlockSize(kxs.np.enc_algo_client_to_server);
            int mac_sc_key_len = MAC.getKeyLen(kxs.np.mac_algo_server_to_client);
            int enc_sc_key_len = BlockCipherFactory.getKeySize(kxs.np.enc_algo_server_to_client);
            int enc_sc_block_len = BlockCipherFactory.getBlockSize(kxs.np.enc_algo_server_to_client);
            km = KeyMaterial.create("SHA1", kxs.H, kxs.K, sessionId, enc_cs_key_len, enc_cs_block_len, mac_cs_key_len,
                                    enc_sc_key_len, enc_sc_block_len, mac_sc_key_len);
        }
        catch (IllegalArgumentException e) {
            return false;
        }

        return true;
    }

    protected void finishKex(boolean clientMode) throws IOException {
        if (sessionId == null) {
            sessionId = kxs.H;
        }

        establishKeyMaterial();
        /* Tell the other side that we start using the new material */
        PacketNewKeys ign = new PacketNewKeys();
        tm.sendKexMessage(ign.getPayload());
        BlockCipher cbc;
        MAC mac;
        Compressor comp;

        try {
            cbc = BlockCipherFactory.createCipher(clientMode ? kxs.np.enc_algo_client_to_server
                                                  : kxs.np.enc_algo_server_to_client, true, clientMode ? km.enc_key_client_to_server
                                                  : km.enc_key_server_to_client, clientMode ? km.initial_iv_client_to_server
                                                  : km.initial_iv_server_to_client);

            try {
                mac = new MAC(clientMode ? kxs.np.mac_algo_client_to_server : kxs.np.mac_algo_server_to_client, clientMode
                              ? km.integrity_key_client_to_server : km.integrity_key_server_to_client);
            }
            catch (DigestException e) {
                throw new IOException(e);
            }

            comp = CompressionFactory.createCompressor(kxs.np.comp_algo_client_to_server);
        }
        catch (IllegalArgumentException f) {
            throw new IOException(String.format("Fatal error initializing ciphers. %s", f.getMessage()));
        }

        tm.changeSendCipher(cbc, mac);
        tm.changeSendCompression(comp);
        tm.kexFinished();
    }

    public static String[] getDefaultServerHostkeyAlgorithmList() {
        return HOSTKEY_ALGS.toArray(new String[HOSTKEY_ALGS.size()]);
    }

    public static void checkServerHostkeyAlgorithmsList(String[] algos) {
        for (final String algo : algos) {
            if (!HOSTKEY_ALGS.contains(algo))
                throw new IllegalArgumentException("Unknown server host key algorithm '" + algo + "'");
        }
    }

    public static String[] getDefaultClientKexAlgorithmList() {
        return KEX_ALGS.toArray(new String[KEX_ALGS.size()]);
    }

    public static String[] getDefaultServerKexAlgorithmList() {
        return KEX_ALGS.toArray(new String[KEX_ALGS.size()]);
    }

    public static void checkKexAlgorithmList(String[] algos) {
        for (final String algo : algos) {
            if (!KEX_ALGS.contains(algo))
                throw new IllegalArgumentException("Unknown kex algorithm '" + algo + "'");
        }
    }
}