/*
 * This file is part of the SSH Library
 *
 * Copyright (c) 2018 by Red Hat, Inc.
 *
 * Authors: Jakub Jelen <jjelen@redhat.com>
 *
 * 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"

#define LIBSSH_STATIC

#include "torture.h"
#include "libssh/sftp.h"
#include "libssh/libssh.h"
#include "libssh/priv.h"
#include "libssh/session.h"
#include "libssh/crypto.h"

#include <errno.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <pwd.h>

static int sshd_setup(void **state)
{
    torture_setup_sshd_server(state, false);

    return 0;
}

static int sshd_teardown(void **state)
{
    torture_teardown_sshd_server(state);

    return 0;
}

static int session_setup(void **state)
{
    struct torture_state *s = *state;
    int verbosity = torture_libssh_verbosity();
    struct passwd *pwd;
    bool b = false;
    int rc;

    pwd = getpwnam("bob");
    assert_non_null(pwd);

    rc = setuid(pwd->pw_uid);
    assert_return_code(rc, errno);

    s->ssh.session = ssh_new();
    assert_non_null(s->ssh.session);

    ssh_options_set(s->ssh.session, SSH_OPTIONS_LOG_VERBOSITY, &verbosity);
    ssh_options_set(s->ssh.session, SSH_OPTIONS_HOST, TORTURE_SSH_SERVER);

    /* Authenticate as alice with bob's pubkey */
    rc = ssh_options_set(s->ssh.session, SSH_OPTIONS_USER, TORTURE_SSH_USER_ALICE);
    assert_int_equal(rc, SSH_OK);

    /* Make sure no other configuration options from system will get used */
    rc = ssh_options_set(s->ssh.session, SSH_OPTIONS_PROCESS_CONFIG, &b);
    assert_ssh_return_code(s->ssh.session, rc);

    /* Make sure we do not interfere with another ssh-agent */
    unsetenv("SSH_AUTH_SOCK");
    unsetenv("SSH_AGENT_PID");

    return 0;
}

static int session_teardown(void **state)
{
    struct torture_state *s = *state;

    ssh_free(s->ssh.session);

    return 0;
}

/* Check that the default limits for rekeying are enforced.
 * the limits are too high for testsuite to verify so
 * we should be fine with checking the values in internal
 * structures
 */
static void torture_rekey_default(void **state)
{
    struct torture_state *s = *state;
    int rc;
    struct ssh_crypto_struct *c = NULL;

    /* Define preferred ciphers: */
    if (ssh_fips_mode()) {
        /* We do not have any FIPS allowed cipher with different block size */
        rc = ssh_options_set(s->ssh.session, SSH_OPTIONS_CIPHERS_C_S,
                             "aes128-gcm@openssh.com");
    } else {
        /* (out) C->S has 8B block */
        rc = ssh_options_set(s->ssh.session, SSH_OPTIONS_CIPHERS_C_S,
                             "chacha20-poly1305@openssh.com");
    }
    assert_ssh_return_code(s->ssh.session, rc);
    /* (in) S->C has 16B block */
    rc = ssh_options_set(s->ssh.session, SSH_OPTIONS_CIPHERS_S_C,
                         "aes128-cbc");
    assert_ssh_return_code(s->ssh.session, rc);

    rc = ssh_connect(s->ssh.session);
    assert_ssh_return_code(s->ssh.session, rc);

    c = s->ssh.session->current_crypto;
    /* The blocks limit is set correctly */
    /* For S->C (in) we have 16B block => 2**(L/4) blocks */
    assert_int_equal(c->in_cipher->max_blocks,
                     (uint64_t)1 << (2 * c->in_cipher->blocksize));
    if (ssh_fips_mode()) {
        /* We do not have any FIPS allowed cipher with different block size */
        assert_int_equal(c->in_cipher->max_blocks,
                         (uint64_t)1 << (2 * c->in_cipher->blocksize));
    } else {
        /* The C->S (out) we have 8B block => 1 GB limit */
        assert_int_equal(c->out_cipher->max_blocks,
                         ((uint64_t)1 << 30) / c->out_cipher->blocksize);
    }

    ssh_disconnect(s->ssh.session);
}

/* We lower the rekey limits manually and check that the rekey
 * really happens when sending data
 */
static void torture_rekey_send(void **state)
{
    struct torture_state *s = *state;
    int rc;
    char data[256];
    unsigned int i;
    uint64_t bytes = 2048; /* 2KB (more than the authentication phase) */
    struct ssh_crypto_struct *c = NULL;
    unsigned char *secret_hash = NULL;

    rc = ssh_options_set(s->ssh.session, SSH_OPTIONS_REKEY_DATA, &bytes);
    assert_ssh_return_code(s->ssh.session, rc);

    rc = ssh_connect(s->ssh.session);
    assert_ssh_return_code(s->ssh.session, rc);

    /* The blocks limit is set correctly */
    c = s->ssh.session->current_crypto;
    assert_int_equal(c->in_cipher->max_blocks,
                     bytes / c->in_cipher->blocksize);
    assert_int_equal(c->out_cipher->max_blocks,
                     bytes / c->out_cipher->blocksize);
    /* We should have less encrypted packets than transfered (first are not encrypted) */
    assert_true(c->out_cipher->packets < s->ssh.session->send_seq);
    assert_true(c->in_cipher->packets < s->ssh.session->recv_seq);
    /* Copy the initial secret hash = session_id so we know we changed keys later */
    secret_hash = malloc(c->digest_len);
    assert_non_null(secret_hash);
    memcpy(secret_hash, c->secret_hash, c->digest_len);

    /* OpenSSH can not rekey before authentication so authenticate here */
    rc = ssh_userauth_none(s->ssh.session, NULL);
    /* This request should return a SSH_REQUEST_DENIED error */
    if (rc == SSH_ERROR) {
        assert_int_equal(ssh_get_error_code(s->ssh.session), SSH_REQUEST_DENIED);
    }
    rc = ssh_userauth_list(s->ssh.session, NULL);
    assert_true(rc & SSH_AUTH_METHOD_PUBLICKEY);

    rc = ssh_userauth_publickey_auto(s->ssh.session, NULL, NULL);
    assert_int_equal(rc, SSH_AUTH_SUCCESS);

    /* send ignore packets of up to 1KB to trigger rekey */
    memset(data, 0, sizeof(data));
    memset(data, 'A', 128);
    for (i = 0; i < 16; i++) {
        ssh_send_ignore(s->ssh.session, data);
        ssh_handle_packets(s->ssh.session, 50);
    }

    /* The rekey limit was restored in the new crypto to the same value */
    c = s->ssh.session->current_crypto;
    assert_int_equal(c->in_cipher->max_blocks, bytes / c->in_cipher->blocksize);
    assert_int_equal(c->out_cipher->max_blocks, bytes / c->out_cipher->blocksize);
    /* Check that the secret hash is different than initially */
    assert_memory_not_equal(secret_hash, c->secret_hash, c->digest_len);
    free(secret_hash);

    ssh_disconnect(s->ssh.session);
}

#ifdef WITH_SFTP
static void session_setup_sftp(void **state)
{
    struct torture_state *s = *state;
    int rc;

    rc = ssh_connect(s->ssh.session);
    assert_ssh_return_code(s->ssh.session, rc);

    /* OpenSSH can not rekey before authentication so authenticate here */
    rc = ssh_userauth_none(s->ssh.session, NULL);
    /* This request should return a SSH_REQUEST_DENIED error */
    if (rc == SSH_ERROR) {
        assert_int_equal(ssh_get_error_code(s->ssh.session), SSH_REQUEST_DENIED);
    }
    rc = ssh_userauth_list(s->ssh.session, NULL);
    assert_true(rc & SSH_AUTH_METHOD_PUBLICKEY);

    rc = ssh_userauth_publickey_auto(s->ssh.session, NULL, NULL);
    assert_int_equal(rc, SSH_AUTH_SUCCESS);

    /* Initialize SFTP session */
    s->ssh.tsftp = torture_sftp_session(s->ssh.session);
    assert_non_null(s->ssh.tsftp);
}

uint64_t bytes = 2048; /* 2KB */

static int session_setup_sftp_client(void **state)
{
    struct torture_state *s = *state;
    int rc;

    session_setup(state);

    rc = ssh_options_set(s->ssh.session, SSH_OPTIONS_REKEY_DATA, &bytes);
    assert_ssh_return_code(s->ssh.session, rc);

    session_setup_sftp(state);

    return 0;
}

#define MAX_XFER_BUF_SIZE 16384

/* To trigger rekey by receiving data, the easiest thing is probably to
 * use sftp
 */
static void torture_rekey_recv(void **state)
{
    struct torture_state *s = *state;
    struct ssh_crypto_struct *c = NULL;
    unsigned char *secret_hash = NULL;

    char libssh_tmp_file[] = "/tmp/libssh_sftp_test_XXXXXX";
    char buf[MAX_XFER_BUF_SIZE];
    ssize_t bytesread;
    ssize_t byteswritten;
    int fd;
    sftp_file file;
    mode_t mask;

    /* The blocks limit is set correctly */
    c = s->ssh.session->current_crypto;
    assert_int_equal(c->in_cipher->max_blocks, bytes / c->in_cipher->blocksize);
    assert_int_equal(c->out_cipher->max_blocks, bytes / c->out_cipher->blocksize);
    /* We should have less encrypted packets than transfered (first are not encrypted) */
    assert_true(c->out_cipher->packets < s->ssh.session->send_seq);
    assert_true(c->in_cipher->packets < s->ssh.session->recv_seq);
    /* Copy the initial secret hash = session_id so we know we changed keys later */
    secret_hash = malloc(c->digest_len);
    assert_non_null(secret_hash);
    memcpy(secret_hash, c->secret_hash, c->digest_len);

    /* Download a file */
    file = sftp_open(s->ssh.tsftp->sftp, "/usr/bin/ssh", O_RDONLY, 0);
    assert_non_null(file);

    mask = umask(S_IRWXO | S_IRWXG);
    fd = mkstemp(libssh_tmp_file);
    umask(mask);
    unlink(libssh_tmp_file);

    for (;;) {
        bytesread = sftp_read(file, buf, MAX_XFER_BUF_SIZE);
        if (bytesread == 0) {
                break; /* EOF */
        }
        assert_false(bytesread < 0);

        byteswritten = write(fd, buf, bytesread);
        assert_int_equal(byteswritten, bytesread);
    }

    close(fd);

    /* The rekey limit was restored in the new crypto to the same value */
    c = s->ssh.session->current_crypto;
    assert_int_equal(c->in_cipher->max_blocks, bytes / c->in_cipher->blocksize);
    assert_int_equal(c->out_cipher->max_blocks, bytes / c->out_cipher->blocksize);
    /* Check that the secret hash is different than initially */
    assert_memory_not_equal(secret_hash, c->secret_hash, c->digest_len);
    free(secret_hash);

    torture_sftp_close(s->ssh.tsftp);
    ssh_disconnect(s->ssh.session);
}
#endif /* WITH_SFTP */

/* Rekey time requires rekey after specified time and is off by default.
 * Setting the time to small enough value and waiting, we should trigger
 * rekey on the first sent packet afterward.
 */
static void torture_rekey_time(void **state)
{
    struct torture_state *s = *state;
    int rc;
    char data[256];
    unsigned int i;
    uint32_t time = 3; /* 3 seconds */
    struct ssh_crypto_struct *c = NULL;
    unsigned char *secret_hash = NULL;

    rc = ssh_options_set(s->ssh.session, SSH_OPTIONS_REKEY_TIME, &time);
    assert_ssh_return_code(s->ssh.session, rc);
    /* The time is internally stored in microseconds */
    assert_int_equal(time * 1000, s->ssh.session->opts.rekey_time);

    rc = ssh_connect(s->ssh.session);
    assert_ssh_return_code(s->ssh.session, rc);

    /* Copy the initial secret hash = session_id so we know we changed keys later */
    c = s->ssh.session->current_crypto;
    secret_hash = malloc(c->digest_len);
    assert_non_null(secret_hash);
    memcpy(secret_hash, c->secret_hash, c->digest_len);

    /* OpenSSH can not rekey before authentication so authenticate here */
    rc = ssh_userauth_none(s->ssh.session, NULL);
    /* This request should return a SSH_REQUEST_DENIED error */
    if (rc == SSH_ERROR) {
        assert_int_equal(ssh_get_error_code(s->ssh.session), SSH_REQUEST_DENIED);
    }
    rc = ssh_userauth_list(s->ssh.session, NULL);
    assert_true(rc & SSH_AUTH_METHOD_PUBLICKEY);

    rc = ssh_userauth_publickey_auto(s->ssh.session, NULL, NULL);
    assert_int_equal(rc, SSH_AUTH_SUCCESS);

    /* Send some data. This should not trigger rekey yet */
    memset(data, 0, sizeof(data));
    memset(data, 'A', 8);
    for (i = 0; i < 3; i++) {
        ssh_send_ignore(s->ssh.session, data);
        ssh_handle_packets(s->ssh.session, 50);
    }

    /* Check that the secret hash is the same */
    c = s->ssh.session->current_crypto;
    assert_memory_equal(secret_hash, c->secret_hash, c->digest_len);

    /* Wait some more time */
    sleep(3);

    /* send some more data to trigger rekey and handle the
     * key exchange "in background" */
    for (i = 0; i < 8; i++) {
        ssh_send_ignore(s->ssh.session, data);
        ssh_handle_packets(s->ssh.session, 50);
    }

    /* Check that the secret hash is different than initially */
    c = s->ssh.session->current_crypto;
    assert_memory_not_equal(secret_hash, c->secret_hash, c->digest_len);
    free(secret_hash);

    ssh_disconnect(s->ssh.session);
}

/* We lower the rekey limits manually and check that the rekey
 * really happens when sending data
 */
static void torture_rekey_server_send(void **state)
{
    struct torture_state *s = *state;
    int rc;
    char data[256];
    unsigned int i;
    struct ssh_crypto_struct *c = NULL;
    unsigned char *secret_hash = NULL;
    const char *sshd_config = "RekeyLimit 2K none";

    torture_update_sshd_config(state, sshd_config);

    rc = ssh_connect(s->ssh.session);
    assert_ssh_return_code(s->ssh.session, rc);

    /* Copy the initial secret hash = session_id so we know we changed keys later */
    c = s->ssh.session->current_crypto;
    secret_hash = malloc(c->digest_len);
    assert_non_null(secret_hash);
    memcpy(secret_hash, c->secret_hash, c->digest_len);

    /* OpenSSH can not rekey before authentication so authenticate here */
    rc = ssh_userauth_none(s->ssh.session, NULL);
    /* This request should return a SSH_REQUEST_DENIED error */
    if (rc == SSH_ERROR) {
        assert_int_equal(ssh_get_error_code(s->ssh.session), SSH_REQUEST_DENIED);
    }
    rc = ssh_userauth_list(s->ssh.session, NULL);
    assert_true(rc & SSH_AUTH_METHOD_PUBLICKEY);

    rc = ssh_userauth_publickey_auto(s->ssh.session, NULL, NULL);
    assert_int_equal(rc, SSH_AUTH_SUCCESS);

    /* send ignore packets of up to 1KB to trigger rekey */
    memset(data, 0, sizeof(data));
    memset(data, 'A', 128);
    for (i = 0; i < 20; i++) {
        ssh_send_ignore(s->ssh.session, data);
        ssh_handle_packets(s->ssh.session, 50);
    }

    /* Check that the secret hash is different than initially */
    c = s->ssh.session->current_crypto;
    assert_memory_not_equal(secret_hash, c->secret_hash, c->digest_len);
    free(secret_hash);

    ssh_disconnect(s->ssh.session);
}

#ifdef WITH_SFTP
static int session_setup_sftp_server(void **state)
{
    const char *sshd_config = "RekeyLimit 2K none";

    session_setup(state);

    torture_update_sshd_config(state, sshd_config);

    session_setup_sftp(state);

    return 0;
}

static void torture_rekey_server_recv(void **state)
{
    struct torture_state *s = *state;
    struct ssh_crypto_struct *c = NULL;
    unsigned char *secret_hash = NULL;
    char libssh_tmp_file[] = "/tmp/libssh_sftp_test_XXXXXX";
    char buf[MAX_XFER_BUF_SIZE];
    ssize_t bytesread;
    ssize_t byteswritten;
    int fd;
    sftp_file file;
    mode_t mask;

    /* Copy the initial secret hash = session_id so we know we changed keys later */
    c = s->ssh.session->current_crypto;
    secret_hash = malloc(c->digest_len);
    assert_non_null(secret_hash);
    memcpy(secret_hash, c->secret_hash, c->digest_len);

    /* Download a file */
    file = sftp_open(s->ssh.tsftp->sftp, "/usr/bin/ssh", O_RDONLY, 0);
    assert_non_null(file);

    mask = umask(S_IRWXO | S_IRWXG);
    fd = mkstemp(libssh_tmp_file);
    umask(mask);
    unlink(libssh_tmp_file);

    for (;;) {
        bytesread = sftp_read(file, buf, MAX_XFER_BUF_SIZE);
        if (bytesread == 0) {
                break; /* EOF */
        }
        assert_false(bytesread < 0);

        byteswritten = write(fd, buf, bytesread);
        assert_int_equal(byteswritten, bytesread);
    }

    close(fd);

    /* Check that the secret hash is different than initially */
    c = s->ssh.session->current_crypto;
    assert_memory_not_equal(secret_hash, c->secret_hash, c->digest_len);
    free(secret_hash);

    torture_sftp_close(s->ssh.tsftp);
    ssh_disconnect(s->ssh.session);
}
#endif /* WITH_SFTP */


int torture_run_tests(void) {
    int rc;
    struct CMUnitTest tests[] = {
        cmocka_unit_test_setup_teardown(torture_rekey_default,
                                        session_setup,
                                        session_teardown),
        cmocka_unit_test_setup_teardown(torture_rekey_time,
                                        session_setup,
                                        session_teardown),
#ifdef WITH_SFTP
        cmocka_unit_test_setup_teardown(torture_rekey_recv,
                                        session_setup_sftp_client,
                                        session_teardown),
#endif /* WITH_SFTP */
        cmocka_unit_test_setup_teardown(torture_rekey_send,
                                        session_setup,
                                        session_teardown),
        /* Note, that this modifies the sshd_config */
        cmocka_unit_test_setup_teardown(torture_rekey_server_send,
                                        session_setup,
                                        session_teardown),
#ifdef WITH_SFTP
        cmocka_unit_test_setup_teardown(torture_rekey_server_recv,
                                        session_setup_sftp_server,
                                        session_teardown),
#endif /* WITH_SFTP */
        /* TODO verify the two rekey are possible and the states are not broken after rekey */
    };

    ssh_init();

    torture_filter_tests(tests);
    rc = cmocka_run_group_tests(tests, sshd_setup, sshd_teardown);

    ssh_finalize();

    return rc;
}