/*

 * This file is part of the SSH Library

 *

 * Copyright (c) 2015 by Aris Adamantiadis

 *

 * 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/libssh.h"

#include "libssh/crypto.h"

#include "libssh/chacha.h"

#include "libssh/poly1305.h"

#include "libssh/misc.h"

/* size of the keys k1 and k2 as defined in specs */

#define CHACHA20_KEYLEN 32

struct chacha20_poly1305_keysched {

    /* key used for encrypting the length field*/

    struct chacha_ctx k1;

    /* key used for encrypting the packets */

    struct chacha_ctx k2;

};

#pragma pack(push, 1)

struct ssh_packet_header {

    uint32_t length;

    uint8_t payload[];

};

#pragma pack(pop)

static const uint8_t zero_block_counter[8] = {0, 0, 0, 0, 0, 0, 0, 0};

static const uint8_t payload_block_counter[8] = {1, 0, 0, 0, 0, 0, 0, 0};

static int chacha20_set_encrypt_key(struct ssh_cipher_struct *cipher,

                                    void *key,

                                    void *IV)

{

    struct chacha20_poly1305_keysched *sched;

    uint8_t *u8key = key;

    (void)IV;

    if (cipher->chacha20_schedule == NULL) {

        sched = malloc(sizeof *sched);

        if (sched == NULL){

            return -1;

        }

    } else {

        sched = cipher->chacha20_schedule;

    }

    chacha_keysetup(&sched->k2, u8key, CHACHA20_KEYLEN * 8);

    chacha_keysetup(&sched->k1, u8key + CHACHA20_KEYLEN, CHACHA20_KEYLEN * 8);

    cipher->chacha20_schedule = sched;

    return 0;

}

/**

 * @internal

 *

 * @brief encrypts an outgoing packet with chacha20 and authenticate it

 * with poly1305.

 */

static void chacha20_poly1305_aead_encrypt(struct ssh_cipher_struct *cipher,

                                           void *in,

                                           void *out,

                                           size_t len,

                                           uint8_t *tag,

                                           uint64_t seq)

{

    struct ssh_packet_header *in_packet = in, *out_packet = out;

    uint8_t poly1305_ctx[POLY1305_KEYLEN] = {0};

    struct chacha20_poly1305_keysched *keys = cipher->chacha20_schedule;

    seq = htonll(seq);

    /* step 1, prepare the poly1305 key */

    chacha_ivsetup(&keys->k2, (uint8_t *)&seq, zero_block_counter);

    chacha_encrypt_bytes(&keys->k2,

                         poly1305_ctx,

                         poly1305_ctx,

                         POLY1305_KEYLEN);

    /* step 2, encrypt length field */

    chacha_ivsetup(&keys->k1, (uint8_t *)&seq, zero_block_counter);

    chacha_encrypt_bytes(&keys->k1,

                         (uint8_t *)&in_packet->length,

                         (uint8_t *)&out_packet->length,

                         sizeof(uint32_t));

    /* step 3, encrypt packet payload */

    chacha_ivsetup(&keys->k2, (uint8_t *)&seq, payload_block_counter);

    chacha_encrypt_bytes(&keys->k2,

                         in_packet->payload,

                         out_packet->payload,

                         len - sizeof(uint32_t));

    /* ssh_print_hexa("poly1305_ctx", poly1305_ctx, sizeof(poly1305_ctx)); */

    /* step 4, compute the MAC */

    poly1305_auth(tag, (uint8_t *)out_packet, len, poly1305_ctx);

    /* ssh_print_hexa("poly1305 src", (uint8_t *)out_packet, len);

    ssh_print_hexa("poly1305 tag", tag, POLY1305_TAGLEN); */

}

static int chacha20_poly1305_aead_decrypt_length(

        struct ssh_cipher_struct *cipher,

        void *in,

        uint8_t *out,

        size_t len,

        uint64_t seq)

{

    struct chacha20_poly1305_keysched *keys = cipher->chacha20_schedule;

    if (len < sizeof(uint32_t)) {

        return SSH_ERROR;

    }

    seq = htonll(seq);

    chacha_ivsetup(&keys->k1, (uint8_t *)&seq, zero_block_counter);

    chacha_encrypt_bytes(&keys->k1,

                         in,

                         (uint8_t *)out,

                         sizeof(uint32_t));

    return SSH_OK;

}

static int chacha20_poly1305_aead_decrypt(struct ssh_cipher_struct *cipher,

                                          void *complete_packet,

                                          uint8_t *out,

                                          size_t encrypted_size,

                                          uint64_t seq)

{

    uint8_t poly1305_ctx[POLY1305_KEYLEN] = {0};

    uint8_t tag[POLY1305_TAGLEN] = {0};

    struct chacha20_poly1305_keysched *keys = cipher->chacha20_schedule;

    uint8_t *mac = (uint8_t *)complete_packet + sizeof(uint32_t) + encrypted_size;

    int cmp;

    seq = htonll(seq);

    ZERO_STRUCT(poly1305_ctx);

    chacha_ivsetup(&keys->k2, (uint8_t *)&seq, zero_block_counter);

    chacha_encrypt_bytes(&keys->k2,

                         poly1305_ctx,

                         poly1305_ctx,

                         POLY1305_KEYLEN);

#if 0

    ssh_print_hexa("poly1305_ctx", poly1305_ctx, sizeof(poly1305_ctx));

#endif

    poly1305_auth(tag, (uint8_t *)complete_packet, encrypted_size +

            sizeof(uint32_t), poly1305_ctx);

#if 0

    ssh_print_hexa("poly1305 src",

                   (uint8_t*)complete_packet,

                   encrypted_size + 4);

    ssh_print_hexa("poly1305 tag", tag, POLY1305_TAGLEN);

    ssh_print_hexa("received tag", mac, POLY1305_TAGLEN);

#endif

    cmp = memcmp(tag, mac, POLY1305_TAGLEN);

    if(cmp != 0) {

        /* mac error */

        SSH_LOG(SSH_LOG_PACKET,"poly1305 verify error");

        return SSH_ERROR;

    }

    chacha_ivsetup(&keys->k2, (uint8_t *)&seq, payload_block_counter);

    chacha_encrypt_bytes(&keys->k2,

                         (uint8_t *)complete_packet + sizeof(uint32_t),

                         out,

                         encrypted_size);

    return SSH_OK;

}

static void chacha20_cleanup(struct ssh_cipher_struct *cipher) {

    SAFE_FREE(cipher->chacha20_schedule);

}

const struct ssh_cipher_struct chacha20poly1305_cipher = {

    .ciphertype = SSH_AEAD_CHACHA20_POLY1305,

    .name = "chacha20-poly1305@openssh.com",

    .blocksize = 8,

    .lenfield_blocksize = 4,

    .keylen = sizeof(struct chacha20_poly1305_keysched),

    .keysize = 512,

    .tag_size = POLY1305_TAGLEN,

    .set_encrypt_key = chacha20_set_encrypt_key,

    .set_decrypt_key = chacha20_set_encrypt_key,

    .aead_encrypt = chacha20_poly1305_aead_encrypt,

    .aead_decrypt_length = chacha20_poly1305_aead_decrypt_length,

    .aead_decrypt = chacha20_poly1305_aead_decrypt,

    .cleanup = chacha20_cleanup

};

const struct ssh_cipher_struct *ssh_get_chacha20poly1305_cipher(void)

{

    return &chacha20poly1305_cipher;

}