/*

 * dh-gex.c - diffie-hellman group exchange

 *

 * This file is part of the SSH Library

 *

 * Copyright (c) 2016 by Aris Adamantiadis <aris@0xbadc0de.be>

 *

 * The SSH Library is free software; you can redistribute it and/or modify

 * it under the terms of the GNU Lesser General Public License as published by

 * the Free Software Foundation; either version 2.1 of the License, or (at your

 * option) any later version.

 *

 * The SSH Library 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 Lesser General Public

 * License for more details.

 *

 * You should have received a copy of the GNU Lesser General Public License

 * along with the SSH Library; see the file COPYING.  If not, write to

 * the Free Software Foundation, Inc., 59 Temple Place - Suite 330, Boston,

 * MA 02111-1307, USA.

 */

#include "config.h"

#include <errno.h>

#include <stdbool.h>

#include <string.h>

#include <stdio.h>

#include "libssh/priv.h"

#include "libssh/dh-gex.h"

#include "libssh/libssh.h"

#include "libssh/ssh2.h"

#include "libssh/callbacks.h"

#include "libssh/dh.h"

#include "libssh/buffer.h"

#include "libssh/session.h"

/* Minimum, recommanded and maximum size of DH group */

#define DH_PMIN 2048

#define DH_PREQ 2048

#define DH_PMAX 8192

static SSH_PACKET_CALLBACK(ssh_packet_client_dhgex_group);

static SSH_PACKET_CALLBACK(ssh_packet_client_dhgex_reply);

static ssh_packet_callback dhgex_client_callbacks[] = {

    ssh_packet_client_dhgex_group, /* SSH_MSG_KEX_DH_GEX_GROUP */

    NULL,                          /* SSH_MSG_KEX_DH_GEX_INIT */

    ssh_packet_client_dhgex_reply  /* SSH_MSG_KEX_DH_GEX_REPLY */

};

static struct ssh_packet_callbacks_struct ssh_dhgex_client_callbacks = {

    .start = SSH2_MSG_KEX_DH_GEX_GROUP,

    .n_callbacks = 3,

    .callbacks = dhgex_client_callbacks,

    .user = NULL

};

/** @internal

 * @brief initiates a diffie-hellman-group-exchange kex

 */

int ssh_client_dhgex_init(ssh_session session)

{

    int rc;

    rc = ssh_dh_init_common(session->next_crypto);

    if (rc != SSH_OK){

        goto error;

    }

    session->next_crypto->dh_pmin = DH_PMIN;

    session->next_crypto->dh_pn = DH_PREQ;

    session->next_crypto->dh_pmax = DH_PMAX;

    /* Minimum group size, preferred group size, maximum group size */

    rc = ssh_buffer_pack(session->out_buffer,

                         "bddd",

                         SSH2_MSG_KEX_DH_GEX_REQUEST,

                         session->next_crypto->dh_pmin,

                         session->next_crypto->dh_pn,

                         session->next_crypto->dh_pmax);

    if (rc != SSH_OK) {

        goto error;

    }

    /* register the packet callbacks */

    ssh_packet_set_callbacks(session, &ssh_dhgex_client_callbacks);

    session->dh_handshake_state = DH_STATE_REQUEST_SENT;

    rc = ssh_packet_send(session);

    if (rc == SSH_ERROR) {

        goto error;

    }

    return rc;

error:

    ssh_dh_cleanup(session->next_crypto);

    return SSH_ERROR;

}

/** @internal

 *  @brief handle a DH_GEX_GROUP packet, client side. This packet contains

 *         the group parameters.

 */

SSH_PACKET_CALLBACK(ssh_packet_client_dhgex_group)

{

    int rc;

    int blen;

    bignum pmin1 = NULL, one = NULL;

    bignum_CTX ctx = bignum_ctx_new();

    bignum modulus, generator;

    const_bignum pubkey;

    (void) type;

    (void) user;

    SSH_LOG(SSH_LOG_PROTOCOL, "SSH_MSG_KEX_DH_GEX_GROUP received");

    if (bignum_ctx_invalid(ctx)) {

        goto error;

    }

    if (session->dh_handshake_state != DH_STATE_REQUEST_SENT) {

        ssh_set_error(session,

                      SSH_FATAL,

                      "Received DH_GEX_GROUP in invalid state");

        goto error;

    }

    one = bignum_new();

    pmin1 = bignum_new();

    if (one == NULL || pmin1 == NULL) {

        ssh_set_error_oom(session);

        goto error;

    }

    rc = ssh_buffer_unpack(packet,

                           "BB",

                           &modulus,

                           &generator);

    if (rc != SSH_OK) {

        ssh_set_error(session, SSH_FATAL, "Invalid DH_GEX_GROUP packet");

        goto error;

    }

    /* basic checks */

    if (ssh_fips_mode() &&

        !ssh_dh_is_known_group(modulus, generator)) {

        ssh_set_error(session,

                      SSH_FATAL,

                      "The received DH group is not FIPS approved");

        goto error;

    }

    rc = bignum_set_word(one, 1);

    if (rc != 1) {

        goto error;

    }

    blen = bignum_num_bits(modulus);

    if (blen < DH_PMIN || blen > DH_PMAX) {

        ssh_set_error(session,

                SSH_FATAL,

                "Invalid dh group parameter p: %d not in [%d:%d]",

                blen,

                DH_PMIN,

                DH_PMAX);

        goto error;

    }

    if (bignum_cmp(modulus, one) <= 0) {

        /* p must be positive and preferably bigger than one */

        ssh_set_error(session, SSH_FATAL, "Invalid dh group parameter p");

    }

    if (!bignum_is_bit_set(modulus, 0)) {

        /* p must be a prime and therefore not divisible by 2 */

        ssh_set_error(session, SSH_FATAL, "Invalid dh group parameter p");

        goto error;

    }

    bignum_sub(pmin1, modulus, one);

    if (bignum_cmp(generator, one) <= 0 ||

        bignum_cmp(generator, pmin1) > 0) {

        /* generator must be at least 2 and smaller than p-1*/

        ssh_set_error(session, SSH_FATAL, "Invalid dh group parameter g");

        goto error;

    }

    bignum_ctx_free(ctx);

    ctx = NULL;

    /* all checks passed, set parameters */

    rc = ssh_dh_set_parameters(session->next_crypto->dh_ctx,

                               modulus, generator);

    if (rc != SSH_OK) {

        bignum_safe_free(modulus);

        bignum_safe_free(generator);

        goto error;

    }

    /* compute and send DH public parameter */

    rc = ssh_dh_keypair_gen_keys(session->next_crypto->dh_ctx,

                                 DH_CLIENT_KEYPAIR);

    if (rc == SSH_ERROR) {

        goto error;

    }

    rc = ssh_dh_keypair_get_keys(session->next_crypto->dh_ctx,

                                 DH_CLIENT_KEYPAIR, NULL, &pubkey);

    if (rc != SSH_OK) {

        goto error;

    }

    rc = ssh_buffer_pack(session->out_buffer,

                         "bB",

                         SSH2_MSG_KEX_DH_GEX_INIT,

                         pubkey);

    if (rc != SSH_OK) {

        goto error;

    }

    session->dh_handshake_state = DH_STATE_INIT_SENT;

    rc = ssh_packet_send(session);

    if (rc == SSH_ERROR) {

        goto error;

    }

    bignum_safe_free(one);

    bignum_safe_free(pmin1);

    return SSH_PACKET_USED;

error:

    bignum_safe_free(one);

    bignum_safe_free(pmin1);

    if(!bignum_ctx_invalid(ctx)) {

        bignum_ctx_free(ctx);

    }

    ssh_dh_cleanup(session->next_crypto);

    session->session_state = SSH_SESSION_STATE_ERROR;

    return SSH_PACKET_USED;

}

static SSH_PACKET_CALLBACK(ssh_packet_client_dhgex_reply)

{

    struct ssh_crypto_struct *crypto=session->next_crypto;

    int rc;

    ssh_string pubkey_blob = NULL;

    bignum server_pubkey = NULL;

    (void)type;

    (void)user;

    SSH_LOG(SSH_LOG_PROTOCOL, "SSH_MSG_KEX_DH_GEX_REPLY received");

    ssh_packet_remove_callbacks(session, &ssh_dhgex_client_callbacks);

    rc = ssh_buffer_unpack(packet,

                           "SBS",

                           &pubkey_blob, &server_pubkey,

                           &crypto->dh_server_signature);

    if (rc == SSH_ERROR) {

        ssh_set_error(session, SSH_FATAL, "Invalid DH_GEX_REPLY packet");

        goto error;

    }

    rc = ssh_dh_keypair_set_keys(crypto->dh_ctx, DH_SERVER_KEYPAIR,

                                 NULL, server_pubkey);

    if (rc != SSH_OK) {

        bignum_safe_free(server_pubkey);

        goto error;

    }

    rc = ssh_dh_import_next_pubkey_blob(session, pubkey_blob);

    ssh_string_free(pubkey_blob);

    if (rc != 0) {

        goto error;

    }

    rc = ssh_dh_compute_shared_secret(session->next_crypto->dh_ctx,

                                      DH_CLIENT_KEYPAIR, DH_SERVER_KEYPAIR,

                                      &session->next_crypto->shared_secret);

    if (rc == SSH_ERROR) {

        ssh_set_error(session, SSH_FATAL, "Could not generate shared secret");

        goto error;

    }

    /* Send the MSG_NEWKEYS */

    if (ssh_buffer_add_u8(session->out_buffer, SSH2_MSG_NEWKEYS) < 0) {

        goto error;

    }

    rc = ssh_packet_send(session);

    if (rc == SSH_ERROR) {

        goto error;

    }

    SSH_LOG(SSH_LOG_PROTOCOL, "SSH_MSG_NEWKEYS sent");

    session->dh_handshake_state = DH_STATE_NEWKEYS_SENT;

    return SSH_PACKET_USED;

error:

    ssh_dh_cleanup(session->next_crypto);

    session->session_state = SSH_SESSION_STATE_ERROR;

    return SSH_PACKET_USED;

}

#ifdef WITH_SERVER

#define MODULI_FILE "/etc/ssh/moduli"

/* 2     "Safe" prime; (p-1)/2 is also prime. */

#define SAFE_PRIME 2

/* 0x04  Probabilistic Miller-Rabin primality tests. */

#define PRIM_TEST_REQUIRED 0x04

/**

 * @internal

 *

 * @brief Determines if the proposed modulus size is more appropriate than the

 * current one.

 *

 * @returns 1 if it's more appropriate. Returns 0 if same or less appropriate

 */

static bool dhgroup_better_size(uint32_t pmin,

                                uint32_t pn,

                                uint32_t pmax,

                                size_t current_size,

                                size_t proposed_size)

{

    if (current_size == proposed_size) {

        return false;

    }

    if (current_size == pn) {

        /* can't do better */

        return false;

    }

    if (current_size == 0 && proposed_size >= pmin && proposed_size <= pmax) {

        return true;

    }

    if (proposed_size < pmin || proposed_size > pmax) {

        /* out of bounds */

        return false;

    }

    if (current_size == 0) {

        /* not in the allowed window */

        return false;

    }

    if (proposed_size >= pn && proposed_size < current_size) {

        return true;

    }

    if (proposed_size <= pn && proposed_size > current_size) {

        return true;

    }

    if (proposed_size >= pn && current_size < pn) {

        return true;

    }

    /* We're in the allowed window but a better match already exists. */

    return false;

}

/** @internal

 * @brief returns 1 with 1/n probability

 * @returns 1 on with P(1/n), 0 with P(n-1/n).

 */

static bool invn_chance(int n)

{

    uint32_t nounce;

    ssh_get_random(&nounce, sizeof(nounce), 0);

    return (nounce % n) == 0;

}

/** @internal

 * @brief retrieves a DH group from an open moduli file.

 */

static int ssh_retrieve_dhgroup_file(FILE *moduli,

                                     uint32_t pmin,

                                     uint32_t pn,

                                     uint32_t pmax,

                                     size_t *best_size,

                                     char **best_generator,

                                     char **best_modulus)

{

    char timestamp[32] = {0};

    char generator[32] = {0};

    char modulus[4096] = {0};

    size_t type, tests, tries, size, proposed_size;

    int firstbyte;

    int rc;

    size_t line = 0;

    size_t best_nlines = 0;

    for(;;) {

        line++;

        firstbyte = getc(moduli);

        if (firstbyte == '#'){

            do {

                firstbyte = getc(moduli);

            } while(firstbyte != '\n' && firstbyte != EOF);

            continue;

        }

        if (firstbyte == EOF) {

            break;

        }

        ungetc(firstbyte, moduli);

        rc = fscanf(moduli,

                    "%31s %zu %zu %zu %zu %31s %4095s\n",

                    timestamp,

                    &type,

                    &tests,

                    &tries,

                    &size,

                    generator,

                    modulus);

        if (rc != 7){

            if (rc == EOF) {

                break;

            }

            SSH_LOG(SSH_LOG_INFO, "Invalid moduli entry line %zu", line);

            do {

                firstbyte = getc(moduli);

            } while(firstbyte != '\n' && firstbyte != EOF);

            continue;

        }

        /* we only want safe primes that were tested */

        if (type != SAFE_PRIME || !(tests & PRIM_TEST_REQUIRED)) {

            continue;

        }

        proposed_size = size + 1;

        if (proposed_size != *best_size &&

            dhgroup_better_size(pmin, pn, pmax, *best_size, proposed_size)) {

            best_nlines = 0;

            *best_size = proposed_size;

        }

        if (proposed_size == *best_size) {

            best_nlines++;

        }

        /* Use reservoir sampling algorithm */

        if (proposed_size == *best_size && invn_chance(best_nlines)) {

            SAFE_FREE(*best_generator);

            SAFE_FREE(*best_modulus);

            *best_generator = strdup(generator);

            if (*best_generator == NULL) {

                return SSH_ERROR;

            }

            *best_modulus = strdup(modulus);

            if (*best_modulus == NULL) {

                SAFE_FREE(*best_generator);

                return SSH_ERROR;

            }

        }

    }

    if (*best_size != 0) {

        SSH_LOG(SSH_LOG_INFO,

                "Selected %zu bits modulus out of %zu candidates in %zu lines",

                *best_size,

                best_nlines - 1,

                line);

    } else {

        SSH_LOG(SSH_LOG_WARNING,

                "No moduli found for [%u:%u:%u]",

                pmin,

                pn,

                pmax);

    }

    return SSH_OK;

}

/** @internal

 * @brief retrieves a DH group from the moduli file based on bits len parameters

 * @param[in] pmin minimum group size in bits

 * @param[in] pn preferred group size

 * @param[in] pmax maximum group size

 * @param[out] size size of the chosen modulus

 * @param[out] p modulus

 * @param[out] g generator

 * @return SSH_OK on success, SSH_ERROR otherwise.

 */

static int ssh_retrieve_dhgroup(uint32_t pmin,

                                uint32_t pn,

                                uint32_t pmax,

                                size_t *size,

                                bignum *p,

                                bignum *g)

{

    FILE *moduli = NULL;

    char *generator = NULL;

    char *modulus = NULL;

    int rc;

    /* In FIPS mode, we can not negotiate arbitrary primes,

     * but just the approved ones */

    if (ssh_fips_mode()) {

        SSH_LOG(SSH_LOG_TRACE, "In FIPS mode, using built-in primes");

        return ssh_fallback_group(pmax, p, g);

    }

    moduli = fopen(MODULI_FILE, "r");

    if (moduli == NULL) {

        SSH_LOG(SSH_LOG_WARNING,

                "Unable to open moduli file: %s",

                strerror(errno));

        return ssh_fallback_group(pmax, p, g);

    }

    *size = 0;

    *p = NULL;

    *g = NULL;

    rc = ssh_retrieve_dhgroup_file(moduli,

                                   pmin,

                                   pn,

                                   pmax,

                                   size,

                                   &generator,

                                   &modulus);

    fclose(moduli);

    if (rc == SSH_ERROR || *size == 0) {

        goto error;

    }

    rc = bignum_hex2bn(generator, g);

    if (rc == 0) {

        goto error;

    }

    rc = bignum_hex2bn(modulus, p);

    if (rc == 0) {

        goto error;

    }

    SAFE_FREE(generator);

    SAFE_FREE(modulus);

    return SSH_OK;

error:

    bignum_safe_free(*g);

    bignum_safe_free(*p);

    SAFE_FREE(generator);

    SAFE_FREE(modulus);

    return SSH_ERROR;

}

static SSH_PACKET_CALLBACK(ssh_packet_server_dhgex_request);

static SSH_PACKET_CALLBACK(ssh_packet_server_dhgex_init);

static ssh_packet_callback dhgex_server_callbacks[]= {

    NULL, /* SSH_MSG_KEX_DH_GEX_REQUEST_OLD */

    NULL, /* SSH_MSG_KEX_DH_GEX_GROUP */

    ssh_packet_server_dhgex_init,   /* SSH_MSG_KEX_DH_GEX_INIT */

    NULL,                           /* SSH_MSG_KEX_DH_GEX_REPLY */

    ssh_packet_server_dhgex_request /* SSH_MSG_GEX_DH_GEX_REQUEST */

};

static struct ssh_packet_callbacks_struct ssh_dhgex_server_callbacks = {

    .start = SSH2_MSG_KEX_DH_GEX_REQUEST_OLD,

    .n_callbacks = 5,

    .callbacks = dhgex_server_callbacks,

    .user = NULL

};

/** @internal

 * @brief sets up the diffie-hellman-groupx kex callbacks

 */

void ssh_server_dhgex_init(ssh_session session){

    /* register the packet callbacks */

    ssh_packet_set_callbacks(session, &ssh_dhgex_server_callbacks);

    ssh_dh_init_common(session->next_crypto);

    session->dh_handshake_state = DH_STATE_INIT;

}

static SSH_PACKET_CALLBACK(ssh_packet_server_dhgex_request)

{

    bignum modulus = NULL, generator = NULL;

    uint32_t pmin, pn, pmax;

    size_t size = 0;

    int rc;

    (void) type;

    (void) user;

    if (session->dh_handshake_state != DH_STATE_INIT) {

        ssh_set_error(session,

                      SSH_FATAL,

                      "Received DH_GEX_REQUEST in invalid state");

        goto error;

    }

    /* Minimum group size, preferred group size, maximum group size */

    rc = ssh_buffer_unpack(packet, "ddd", &pmin, &pn, &pmax);

    if (rc != SSH_OK){

        ssh_set_error_invalid(session);

        goto error;

    }

    SSH_LOG(SSH_LOG_INFO, "dh-gex: DHGEX_REQUEST[%u:%u:%u]", pmin, pn, pmax);

    if (pmin > pn || pn > pmax || pn > DH_PMAX || pmax < DH_PMIN) {

        ssh_set_error(session,

                      SSH_FATAL,

                      "Invalid dh-gex arguments [%u:%u:%u]",

                      pmin,

                      pn,

                      pmax);

        goto error;

    }

    session->next_crypto->dh_pmin = pmin;

    session->next_crypto->dh_pn = pn;

    session->next_crypto->dh_pmax = pmax;

    /* ensure safe parameters */

    if (pmin < DH_PMIN) {

        pmin = DH_PMIN;

        if (pn < pmin) {

            pn = pmin;

        }

    }

    rc = ssh_retrieve_dhgroup(pmin,

                              pn,

                              pmax,

                              &size,

                              &modulus,

                              &generator);

    if (rc == SSH_ERROR) {

        ssh_set_error(session,

                      SSH_FATAL,

                      "Couldn't find DH group for [%u:%u:%u]",

                      pmin,

                      pn,

                      pmax);

        goto error;

    }

    rc = ssh_dh_set_parameters(session->next_crypto->dh_ctx,

                               modulus, generator);

    if (rc != SSH_OK) {

        bignum_safe_free(generator);

        bignum_safe_free(modulus);

        goto error;

    }

    rc = ssh_buffer_pack(session->out_buffer,

                         "bBB",

                         SSH2_MSG_KEX_DH_GEX_GROUP,

                         modulus,

                         generator);

#ifdef HAVE_LIBCRYPTO

        bignum_safe_free(generator);

        bignum_safe_free(modulus);

#endif

    if (rc != SSH_OK) {

        ssh_set_error_invalid(session);

        goto error;

    }

    session->dh_handshake_state = DH_STATE_GROUP_SENT;

    rc = ssh_packet_send(session);

    if (rc == SSH_ERROR) {

        goto error;

    }

error:

    return SSH_PACKET_USED;

}

/** @internal

 * @brief parse an incoming SSH_MSG_KEX_DH_GEX_INIT packet and complete

 *        Diffie-Hellman key exchange

 **/

static SSH_PACKET_CALLBACK(ssh_packet_server_dhgex_init){

    (void) type;

    (void) user;

    SSH_LOG(SSH_LOG_DEBUG, "Received SSH_MSG_KEX_DHGEX_INIT");

    ssh_packet_remove_callbacks(session, &ssh_dhgex_server_callbacks);

    ssh_server_dh_process_init(session, packet);

    return SSH_PACKET_USED;

}

#endif /* WITH_SERVER */