/*
 * This file is part of the SSH Library
 *
 * Copyright (c) 2019 by Simo Sorce - Red Hat, Inc.
 *
 * 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 "libssh/session.h"
#include "libssh/dh.h"
#include "libssh/buffer.h"
#include "libssh/ssh2.h"
#include "libssh/pki.h"
#include "libssh/bignum.h"

#include "openssl/crypto.h"
#include "openssl/dh.h"
#include "libcrypto-compat.h"

extern bignum ssh_dh_generator;
extern bignum ssh_dh_group1;
extern bignum ssh_dh_group14;
extern bignum ssh_dh_group16;
extern bignum ssh_dh_group18;

struct dh_ctx {
    DH *keypair[2];
};

void ssh_dh_debug_crypto(struct ssh_crypto_struct *c)
{
#ifdef DEBUG_CRYPTO
    const_bignum x = NULL, y = NULL, e = NULL, f = NULL;

    ssh_dh_keypair_get_keys(c->dh_ctx, DH_CLIENT_KEYPAIR, &x, &e);
    ssh_dh_keypair_get_keys(c->dh_ctx, DH_SERVER_KEYPAIR, &y, &f);
    ssh_print_bignum("x", x);
    ssh_print_bignum("y", y);
    ssh_print_bignum("e", e);
    ssh_print_bignum("f", f);

    ssh_log_hexdump("Session server cookie", c->server_kex.cookie, 16);
    ssh_log_hexdump("Session client cookie", c->client_kex.cookie, 16);
    ssh_print_bignum("k", c->shared_secret);

#else
    (void)c; /* UNUSED_PARAM */
#endif
}

int ssh_dh_keypair_get_keys(struct dh_ctx *ctx, int peer,
                            const_bignum *priv, const_bignum *pub)
{
    if (((peer != DH_CLIENT_KEYPAIR) && (peer != DH_SERVER_KEYPAIR)) ||
        ((priv == NULL) && (pub == NULL)) || (ctx == NULL) ||
        (ctx->keypair[peer] == NULL)) {
        return SSH_ERROR;
    }
    DH_get0_key(ctx->keypair[peer], pub, priv);
    if (priv && (*priv == NULL || bignum_num_bits(*priv) == 0)) {
        return SSH_ERROR;
    }
    if (pub && (*pub == NULL || bignum_num_bits(*pub) == 0)) {
        return SSH_ERROR;
    }

    return SSH_OK;
}

int ssh_dh_keypair_set_keys(struct dh_ctx *ctx, int peer,
                            const bignum priv, const bignum pub)
{
    bignum priv_key = NULL;
    bignum pub_key = NULL;

    if (((peer != DH_CLIENT_KEYPAIR) && (peer != DH_SERVER_KEYPAIR)) ||
        ((priv == NULL) && (pub == NULL)) || (ctx == NULL) ||
        (ctx->keypair[peer] == NULL)) {
        return SSH_ERROR;
    }

    if (priv) {
        priv_key = priv;
    }
    if (pub) {
        pub_key = pub;
    }
    (void)DH_set0_key(ctx->keypair[peer], pub_key, priv_key);

    return SSH_OK;
}

int ssh_dh_get_parameters(struct dh_ctx *ctx,
                          const_bignum *modulus, const_bignum *generator)
{
    if (ctx == NULL || ctx->keypair[0] == NULL) {
        return SSH_ERROR;
    }
    DH_get0_pqg(ctx->keypair[0], modulus, NULL, generator);
    return SSH_OK;
}

int ssh_dh_set_parameters(struct dh_ctx *ctx,
                          const bignum modulus, const bignum generator)
{
    size_t i;
    int rc;

    if ((ctx == NULL) || (modulus == NULL) || (generator == NULL)) {
        return SSH_ERROR;
    }

    for (i = 0; i < 2; i++) {
        bignum p = NULL;
        bignum g = NULL;

        /* when setting modulus or generator,
         * make sure to invalidate existing keys */
        DH_free(ctx->keypair[i]);
        ctx->keypair[i] = DH_new();
        if (ctx->keypair[i] == NULL) {
            rc = SSH_ERROR;
            goto done;
        }

        p = BN_dup(modulus);
        g = BN_dup(generator);
        rc = DH_set0_pqg(ctx->keypair[i], p, NULL, g);
        if (rc != 1) {
            BN_free(p);
            BN_free(g);
            rc = SSH_ERROR;
            goto done;
        }
    }

    rc = SSH_OK;
done:
    if (rc != SSH_OK) {
        DH_free(ctx->keypair[0]);
        DH_free(ctx->keypair[1]);
        ctx->keypair[0] = NULL;
        ctx->keypair[1] = NULL;
    }
    return rc;
}

/**
 * @internal
 * @brief allocate and initialize ephemeral values used in dh kex
 */
int ssh_dh_init_common(struct ssh_crypto_struct *crypto)
{
    struct dh_ctx *ctx;
    int rc;

    ctx = calloc(1, sizeof(*ctx));
    if (ctx == NULL) {
        return SSH_ERROR;
    }
    crypto->dh_ctx = ctx;

    switch (crypto->kex_type) {
    case SSH_KEX_DH_GROUP1_SHA1:
        rc = ssh_dh_set_parameters(ctx, ssh_dh_group1, ssh_dh_generator);
        break;
    case SSH_KEX_DH_GROUP14_SHA1:
    case SSH_KEX_DH_GROUP14_SHA256:
        rc = ssh_dh_set_parameters(ctx, ssh_dh_group14, ssh_dh_generator);
        break;
    case SSH_KEX_DH_GROUP16_SHA512:
        rc = ssh_dh_set_parameters(ctx, ssh_dh_group16, ssh_dh_generator);
        break;
    case SSH_KEX_DH_GROUP18_SHA512:
        rc = ssh_dh_set_parameters(ctx, ssh_dh_group18, ssh_dh_generator);
        break;
    default:
        rc = SSH_OK;
        break;
    }

    if (rc != SSH_OK) {
        ssh_dh_cleanup(crypto);
    }
    return rc;
}

void ssh_dh_cleanup(struct ssh_crypto_struct *crypto)
{
    if (crypto->dh_ctx != NULL) {
        DH_free(crypto->dh_ctx->keypair[0]);
        DH_free(crypto->dh_ctx->keypair[1]);
        free(crypto->dh_ctx);
        crypto->dh_ctx = NULL;
    }
}

/** @internal
 * @brief generates a secret DH parameter of at least DH_SECURITY_BITS
 *        security as well as the corresponding public key.
 * @param[out] parms a dh_ctx that will hold the new keys.
 * @param peer Select either client or server key storage. Valid values are:
 *        DH_CLIENT_KEYPAIR or DH_SERVER_KEYPAIR
 *
 * @return SSH_OK on success, SSH_ERROR on error
 */
int ssh_dh_keypair_gen_keys(struct dh_ctx *dh_ctx, int peer)
{
    int rc;

    if ((dh_ctx == NULL) || (dh_ctx->keypair[peer] == NULL)) {
        return SSH_ERROR;
    }
    rc = DH_generate_key(dh_ctx->keypair[peer]);
    if (rc != 1) {
        return SSH_ERROR;
    }
    return SSH_OK;
}

/** @internal
 * @brief generates a shared secret between the local peer and the remote
 *        peer. The local peer must have been initialized using either the
 *        ssh_dh_keypair_gen_keys() function or by seetting manually both
 *        the private and public keys. The remote peer only needs to have
 *        the remote's peer public key set.
 * @param[in] local peer identifier (DH_CLIENT_KEYPAIR or DH_SERVER_KEYPAIR)
 * @param[in] remote peer identifier (DH_CLIENT_KEYPAIR or DH_SERVER_KEYPAIR)
 * @param[out] dest a new bignum with the shared secret value is returned.
 * @return SSH_OK on success, SSH_ERROR on error
 */
int ssh_dh_compute_shared_secret(struct dh_ctx *dh_ctx, int local, int remote,
                                 bignum *dest)
{
    unsigned char *kstring = NULL;
    const_bignum pub_key = NULL;
    int klen, rc;

    if ((dh_ctx == NULL) ||
        (dh_ctx->keypair[local] == NULL) ||
        (dh_ctx->keypair[remote] == NULL)) {
        return SSH_ERROR;
    }

    kstring = malloc(DH_size(dh_ctx->keypair[local]));
    if (kstring == NULL) {
        rc = SSH_ERROR;
        goto done;
    }

    rc = ssh_dh_keypair_get_keys(dh_ctx, remote, NULL, &pub_key);
    if (rc != SSH_OK) {
        rc = SSH_ERROR;
        goto done;
    }

    klen = DH_compute_key(kstring, pub_key, dh_ctx->keypair[local]);
    if (klen == -1) {
        rc = SSH_ERROR;
        goto done;
    }

    *dest = BN_bin2bn(kstring, klen, NULL);
    if (*dest == NULL) {
        rc = SSH_ERROR;
        goto done;
    }

    rc = SSH_OK;
done:
    free(kstring);
    return rc;
}