* Elliptic Curve Key Exchange
 */

static int curve25519_sha256(LIBSSH2_SESSION *session, unsigned char *data,
                             size_t data_len, unsigned char public_key[LIBSSH2_ED25519_KEY_LEN],
                             unsigned char private_key[LIBSSH2_ED25519_KEY_LEN],
                             kmdhgGPshakex_state_t *exchange_state)
{
    int ret = 0;
    int rc;
    int public_key_len = LIBSSH2_ED25519_KEY_LEN;

    if(data_len < 5) {
        return _libssh2_error(session, LIBSSH2_ERROR_HOSTKEY_INIT,
                              "Data is too short");
    }

    if(exchange_state->state == libssh2_NB_state_idle) {

        /* Setup initial values */
        exchange_state->k = _libssh2_bn_init();

        exchange_state->state = libssh2_NB_state_created;
    }

    if(exchange_state->state == libssh2_NB_state_created) {
        /* parse INIT reply data */
        unsigned char *server_public_key, *server_host_key;
        unsigned int server_public_key_len;
        struct string_buf buf = { .len = 0, .offset = 0 };

        if(data_len < 5) {
            ret = _libssh2_error(session, LIBSSH2_ERROR_PROTO,
                                 "Unexpected key length");
            goto clean_exit;
        }
    
        buf.data = data;
        buf.len = data_len;
        buf.dataptr = buf.data;
        buf.dataptr++; /* advance past packet type */

        if((session->server_hostkey_len = _libssh2_get_c_string(&buf, &server_host_key)) <= 0) {
            ret = _libssh2_error(session, LIBSSH2_ERROR_PROTO,
                                 "Unexpected key length");
            goto clean_exit;
        }
        
        session->server_hostkey = LIBSSH2_ALLOC(session, session->server_hostkey_len);
        if (!session->server_hostkey) {
            ret = _libssh2_error(session, LIBSSH2_ERROR_ALLOC,
                                 "Unable to allocate memory for a copy "
                                 "of the host key");
            goto clean_exit;
        }

        memcpy(session->server_hostkey, server_host_key, session->server_hostkey_len);

#if LIBSSH2_MD5
        {
            libssh2_md5_ctx fingerprint_ctx;

            if(libssh2_md5_init(&fingerprint_ctx)) {
                libssh2_md5_update(fingerprint_ctx, session->server_hostkey,
                                   session->server_hostkey_len);
                libssh2_md5_final(fingerprint_ctx, session->server_hostkey_md5);
                session->server_hostkey_md5_valid = TRUE;
            }
            else {
                session->server_hostkey_md5_valid = FALSE;
            }
        }
#ifdef LIBSSH2DEBUG
        {
            char fingerprint[50], *fprint = fingerprint;
            int i;
            for(i = 0; i < 16; i++, fprint += 3) {
                snprintf(fprint, 4, "%02x:", session->server_hostkey_md5[i]);
            }
            *(--fprint) = '\0';
            _libssh2_debug(session, LIBSSH2_TRACE_KEX,
                             "Server's MD5 Fingerprint: %s", fingerprint);
        }
#endif /* LIBSSH2DEBUG */
#endif /* ! LIBSSH2_MD5 */

        {
            libssh2_sha1_ctx fingerprint_ctx;

            if(libssh2_sha1_init(&fingerprint_ctx)) {
                libssh2_sha1_update(fingerprint_ctx, session->server_hostkey,
                                    session->server_hostkey_len);
                libssh2_sha1_final(fingerprint_ctx, session->server_hostkey_sha1);
                session->server_hostkey_sha1_valid = TRUE;
            }
            else {
                session->server_hostkey_sha1_valid = FALSE;
            }
        }
#ifdef LIBSSH2DEBUG
        {
            char fingerprint[64], *fprint = fingerprint;
            int i;

            for(i = 0; i < 20; i++, fprint += 3) {
                snprintf(fprint, 4, "%02x:", session->server_hostkey_sha1[i]);
            }
            *(--fprint) = '\0';
            _libssh2_debug(session, LIBSSH2_TRACE_KEX,
                             "Server's SHA1 Fingerprint: %s", fingerprint);
        }
#endif /* LIBSSH2DEBUG */

        /* SHA256 */
        {
            libssh2_sha256_ctx fingerprint_ctx;

            if(libssh2_sha256_init(&fingerprint_ctx)) {
                libssh2_sha256_update(fingerprint_ctx, session->server_hostkey,
                                      session->server_hostkey_len);
                libssh2_sha256_final(fingerprint_ctx, session->server_hostkey_sha256);
                session->server_hostkey_sha256_valid = TRUE;
            }
            else {
                session->server_hostkey_sha256_valid = FALSE;
            }
        }
#ifdef LIBSSH2DEBUG
        {
            char *base64Fingerprint = NULL;
            _libssh2_base64_encode(session, (const char*)session->server_hostkey_sha256, SHA256_DIGEST_LENGTH, &base64Fingerprint);
            if( base64Fingerprint != NULL) {
                _libssh2_debug(session, LIBSSH2_TRACE_KEX,
                                  "Server's SHA256 Fingerprint: %s", base64Fingerprint);
                LIBSSH2_FREE(session, base64Fingerprint);
            }
        }
#endif /* LIBSSH2DEBUG */

        if(session->hostkey->init(session, session->server_hostkey,
                                   session->server_hostkey_len,
                                   &session->server_hostkey_abstract)) {
            ret = _libssh2_error(session, LIBSSH2_ERROR_HOSTKEY_INIT,
                                 "Unable to initialize hostkey importer");
            goto clean_exit;
        }

        /* server public key Q_S */
        if((server_public_key_len = 
            _libssh2_get_c_string(&buf, &server_public_key)) <= 0) {
            ret = _libssh2_error(session, LIBSSH2_ERROR_PROTO,
                                     "Unexpected key length");
            goto clean_exit;
        }

        if( server_public_key_len != LIBSSH2_ED25519_KEY_LEN) {
            ret = _libssh2_error(session, LIBSSH2_ERROR_HOSTKEY_INIT,
                                 "Unexpected curve25519 server public key length");
            goto clean_exit;
        }

        /* server signature */
        if((exchange_state->h_sig_len =
            _libssh2_get_c_string(&buf, &exchange_state->h_sig)) <= 0) {
            ret = _libssh2_error(session, LIBSSH2_ERROR_HOSTKEY_INIT,
                                 "Unexpected curve25519 server sig length");
            goto clean_exit;

        }

        // Compute the shared secret K
        rc = _libssh2_curve25519_gen_k(&exchange_state->k, private_key, server_public_key);
        if( rc != 0) {
            ret = _libssh2_error(session, LIBSSH2_ERROR_KEX_FAILURE,
                                 "Unable to create ECDH shared secret");
            goto clean_exit;
        }

        exchange_state->k_value_len = _libssh2_bn_bytes(exchange_state->k) + 5;
        if(_libssh2_bn_bits(exchange_state->k) % 8) {
            // don't need leading 00
            exchange_state->k_value_len--;
        }
        exchange_state->k_value =
        LIBSSH2_ALLOC(session, exchange_state->k_value_len);
        if(!exchange_state->k_value) {
            ret = _libssh2_error(session, LIBSSH2_ERROR_ALLOC,
                                 "Unable to allocate buffer for K");
            goto clean_exit;
        }
        _libssh2_htonu32(exchange_state->k_value,
                         exchange_state->k_value_len - 4);
        if(_libssh2_bn_bits(exchange_state->k) % 8) {
            _libssh2_bn_to_bin(exchange_state->k, exchange_state->k_value + 4);
        } else {
            exchange_state->k_value[4] = 0;
            _libssh2_bn_to_bin(exchange_state->k, exchange_state->k_value + 5);
        }

        // verify hash
        LIBSSH2_KEX_METHOD_EC_SHA_HASH_CREATE_VERIFY(256);

        if( rc != 0) {
            ret = _libssh2_error(session, LIBSSH2_ERROR_HOSTKEY_SIGN,
                                 "Unable to verify hostkey signature");
            goto clean_exit;
        }

        exchange_state->c = SSH_MSG_NEWKEYS;
        exchange_state->state = libssh2_NB_state_sent;
    }

    if(exchange_state->state == libssh2_NB_state_sent) {
        rc = _libssh2_transport_send(session, &exchange_state->c, 1, NULL, 0);
        if(rc == LIBSSH2_ERROR_EAGAIN) {
            return rc;
        } else if(rc) {
            ret = _libssh2_error(session, rc, "Unable to send NEWKEYS message");
            goto clean_exit;
        }

        exchange_state->state = libssh2_NB_state_sent2;
    }

    if(exchange_state->state == libssh2_NB_state_sent2) {
        rc = _libssh2_packet_require(session, SSH_MSG_NEWKEYS,
                                     &exchange_state->tmp,
                                     &exchange_state->tmp_len, 0, NULL, 0,
                                     &exchange_state->req_state);
        if(rc == LIBSSH2_ERROR_EAGAIN) {
            return rc;
        } else if(rc) {
            ret = _libssh2_error(session, rc, "Timed out waiting for NEWKEYS");
            goto clean_exit;
        }

        // The first key exchange has been performed, switch to active crypt/comp/mac mode
        session->state |= LIBSSH2_STATE_NEWKEYS;
        _libssh2_debug(session, LIBSSH2_TRACE_KEX, "Received NEWKEYS message");

        // This will actually end up being just packet_type(1) for this packet type anyway
        LIBSSH2_FREE(session, exchange_state->tmp);

        if(!session->session_id) {

            size_t digest_length = SHA256_DIGEST_LENGTH;
            session->session_id = LIBSSH2_ALLOC(session, digest_length);
            if(!session->session_id) {
                ret = _libssh2_error(session, LIBSSH2_ERROR_ALLOC,
                                     "Unable to allocate buffer for SHA digest");
                goto clean_exit;
            }
            memcpy(session->session_id, exchange_state->h_sig_comp,
                   digest_length);
            session->session_id_len = digest_length;
            _libssh2_debug(session, LIBSSH2_TRACE_KEX, "session_id calculated");
        }

        // Cleanup any existing cipher
        if(session->local.crypt->dtor) {
            session->local.crypt->dtor(session,
                                        &session->local.crypt_abstract);
        }

        // Calculate IV/Secret/Key for each direction
        if(session->local.crypt->init) {
            unsigned char *iv = NULL, *secret = NULL;
            int free_iv = 0, free_secret = 0;

            LIBSSH2_KEX_METHOD_SHA_VALUE_HASH(256, iv,
                                              session->local.crypt->
                                              iv_len, "A");
            if(!iv) {
                ret = -1;
                goto clean_exit;
            }

            LIBSSH2_KEX_METHOD_SHA_VALUE_HASH(256, secret,
                                              session->local.crypt->
                                              secret_len, "C");

            if(!secret) {
                LIBSSH2_FREE(session, iv);
                ret = LIBSSH2_ERROR_KEX_FAILURE;
                goto clean_exit;
            }
            if(session->local.crypt->
                init(session, session->local.crypt, iv, &free_iv, secret,
                     &free_secret, 1, &session->local.crypt_abstract)) {
                    LIBSSH2_FREE(session, iv);
                    LIBSSH2_FREE(session, secret);
                    ret = LIBSSH2_ERROR_KEX_FAILURE;
                    goto clean_exit;
                }

            if(free_iv) {
                _libssh2_explicit_zero(iv, session->local.crypt->iv_len);
                LIBSSH2_FREE(session, iv);
            }

            if(free_secret) {
                _libssh2_explicit_zero(secret, session->local.crypt->secret_len);
                LIBSSH2_FREE(session, secret);
            }
        }
        _libssh2_debug(session, LIBSSH2_TRACE_KEX,
                        "Client to Server IV and Key calculated");

        if(session->remote.crypt->dtor) {
            // Cleanup any existing cipher
            session->remote.crypt->dtor(session,
                                        &session->remote.crypt_abstract);
        }

        if(session->remote.crypt->init) {
            unsigned char *iv = NULL, *secret = NULL;
            int free_iv = 0, free_secret = 0;

            LIBSSH2_KEX_METHOD_SHA_VALUE_HASH(256, iv,
                                              session->remote.crypt->
                                              iv_len, "B");

            if(!iv) {
                ret = LIBSSH2_ERROR_KEX_FAILURE;
                goto clean_exit;
            }
            LIBSSH2_KEX_METHOD_SHA_VALUE_HASH(256, secret,
                                              session->remote.crypt->
                                              secret_len, "D");

            if(!secret) {
                LIBSSH2_FREE(session, iv);
                ret = LIBSSH2_ERROR_KEX_FAILURE;
                goto clean_exit;
            }
            if(session->remote.crypt->
                init(session, session->remote.crypt, iv, &free_iv, secret,
                     &free_secret, 0, &session->remote.crypt_abstract)) {
                    LIBSSH2_FREE(session, iv);
                    LIBSSH2_FREE(session, secret);
                    ret = LIBSSH2_ERROR_KEX_FAILURE;
                    goto clean_exit;
                }

            if(free_iv) {
                _libssh2_explicit_zero(iv, session->remote.crypt->iv_len);
                LIBSSH2_FREE(session, iv);
            }

            if(free_secret) {
                _libssh2_explicit_zero(secret, session->remote.crypt->secret_len);
                LIBSSH2_FREE(session, secret);
            }
        }
        _libssh2_debug(session, LIBSSH2_TRACE_KEX,
                        "Server to Client IV and Key calculated");

        if(session->local.mac->dtor) {
            session->local.mac->dtor(session, &session->local.mac_abstract);
        }

        if(session->local.mac->init) {
            unsigned char *key = NULL;
            int free_key = 0;

            LIBSSH2_KEX_METHOD_SHA_VALUE_HASH(256, key,
                                              session->local.mac->
                                              key_len, "E");

            if(!key) {
                ret = LIBSSH2_ERROR_KEX_FAILURE;
                goto clean_exit;
            }
            session->local.mac->init(session, key, &free_key,
                                     &session->local.mac_abstract);

            if(free_key) {
                _libssh2_explicit_zero(key, session->local.mac->key_len);
                LIBSSH2_FREE(session, key);
            }
        }
        _libssh2_debug(session, LIBSSH2_TRACE_KEX,
                        "Client to Server HMAC Key calculated");

        if(session->remote.mac->dtor) {
            session->remote.mac->dtor(session, &session->remote.mac_abstract);
        }

        if(session->remote.mac->init) {
            unsigned char *key = NULL;
            int free_key = 0;

            LIBSSH2_KEX_METHOD_SHA_VALUE_HASH(256, key,
                                              session->remote.mac->
                                              key_len, "F");

            if(!key) {
                ret = LIBSSH2_ERROR_KEX_FAILURE;
                goto clean_exit;
            }
            session->remote.mac->init(session, key, &free_key,
                                      &session->remote.mac_abstract);

            if(free_key) {
                _libssh2_explicit_zero(key, session->remote.mac->key_len);
                LIBSSH2_FREE(session, key);
            }
        }
        _libssh2_debug(session, LIBSSH2_TRACE_KEX,
                        "Server to Client HMAC Key calculated");

        // Initialize compression for each direction

        // Cleanup any existing compression
        if(session->local.comp && session->local.comp->dtor) {
            session->local.comp->dtor(session, 1,
                                      &session->local.comp_abstract);
        }

        if(session->local.comp && session->local.comp->init) {
            if(session->local.comp->init(session, 1,
                                            &session->local.comp_abstract)) {
                ret = LIBSSH2_ERROR_KEX_FAILURE;
                goto clean_exit;
            }
        }
        _libssh2_debug(session, LIBSSH2_TRACE_KEX,
                        "Client to Server compression initialized");

        if(session->remote.comp && session->remote.comp->dtor) {
            session->remote.comp->dtor(session, 0,
                                        &session->remote.comp_abstract);
        }

        if(session->remote.comp && session->remote.comp->init) {
            if(session->remote.comp->init(session, 0,
                                             &session->remote.comp_abstract)) {
                ret = LIBSSH2_ERROR_KEX_FAILURE;
                goto clean_exit;
            }
        }
        _libssh2_debug(session, LIBSSH2_TRACE_KEX,
                        "Server to Client compression initialized");
    }

clean_exit:
    _libssh2_bn_free(exchange_state->k);
    exchange_state->k = NULL;

    if(exchange_state->k_value) {
        LIBSSH2_FREE(session, exchange_state->k_value);
        exchange_state->k_value = NULL;
    }

    exchange_state->state = libssh2_NB_state_idle;