/*
 * Slitheen - a decoy routing system for censorship resistance
 * Copyright (C) 2017 Cecylia Bocovich (cbocovic@uwaterloo.ca)
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, version 3.
 *
 * This program 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 General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 * Additional permission under GNU GPL version 3 section 7
 *
 * If you modify this Program, or any covered work, by linking or combining
 * it with the OpenSSL library (or a modified version of that library),
 * containing parts covered by the terms of the OpenSSL Licence and the
 * SSLeay license, the licensors of this Program grant you additional
 * permission to convey the resulting work. Corresponding Source for a
 * non-source form of such a combination shall include the source code
 * for the parts of the OpenSSL library used as well as that of the covered
 * work.
 */

#include <string.h>

#include <openssl/evp.h>
#include <openssl/dh.h>
#include <openssl/bn.h>
#include <openssl/err.h>
#include <openssl/rand.h>
#include <openssl/ssl.h>

#include <netinet/in.h>

#include "crypto.h"
#include "socks5proxy.h"
#include "tagging.h"
#include "ptwist.h"
#include "util.h"

static super_data *super;

/* PRF using sha384, as defined in RFC 5246 */
int PRF(uint8_t *secret, int32_t secret_len,
        uint8_t *seed1, int32_t seed1_len,
        uint8_t *seed2, int32_t seed2_len,
        uint8_t *seed3, int32_t seed3_len,
        uint8_t *seed4, int32_t seed4_len,
        uint8_t *output, int32_t output_len){

    EVP_MD_CTX *ctx, *ctx_tmp, *ctx_init;
    EVP_PKEY *mac_key;
    const EVP_MD *md = EVP_sha256();

    uint8_t A[EVP_MAX_MD_SIZE];
    size_t len, A_len;
    int chunk = EVP_MD_size(md);
    int remaining = output_len;

    uint8_t *out = output;

#if OPENSSL_VERSION_NUMBER >= 0x1010000eL
    ctx = EVP_MD_CTX_new();
    ctx_tmp = EVP_MD_CTX_new();
    ctx_init = EVP_MD_CTX_new();
#else
    ctx = ecalloc(1, sizeof(EVP_MD_CTX));
    EVP_MD_CTX_init(ctx);
    ctx_tmp = ecalloc(1, sizeof(EVP_MD_CTX));
    EVP_MD_CTX_init(ctx_tmp);
    ctx_init = ecalloc(1, sizeof(EVP_MD_CTX));
    EVP_MD_CTX_init(ctx_init);
#endif

    EVP_MD_CTX_set_flags(ctx_init, EVP_MD_CTX_FLAG_NON_FIPS_ALLOW);

    mac_key = EVP_PKEY_new_mac_key(EVP_PKEY_HMAC, NULL, secret, secret_len);

    /* Calculate first A value */
    EVP_DigestSignInit(ctx_init, NULL, md, NULL, mac_key);
    EVP_MD_CTX_copy_ex(ctx, ctx_init);
    if(seed1 != NULL && seed1_len > 0){
        EVP_DigestSignUpdate(ctx, seed1, seed1_len);
    }
    if(seed2 != NULL && seed2_len > 0){
        EVP_DigestSignUpdate(ctx, seed2, seed2_len);
    }
    if(seed3 != NULL && seed3_len > 0){
        EVP_DigestSignUpdate(ctx, seed3, seed3_len);
    }
    if(seed4 != NULL && seed4_len > 0){
        EVP_DigestSignUpdate(ctx, seed4, seed4_len);
    }
    EVP_DigestSignFinal(ctx, A, &A_len);

    //iterate until desired length is achieved
    while(remaining > 0){
        /* Now compute SHA384(secret, A+seed) */
        EVP_MD_CTX_copy_ex(ctx, ctx_init);
        EVP_DigestSignUpdate(ctx, A, A_len);
        EVP_MD_CTX_copy_ex(ctx_tmp, ctx);
        if(seed1 != NULL && seed1_len > 0){
            EVP_DigestSignUpdate(ctx, seed1, seed1_len);
        }
        if(seed2 != NULL && seed2_len > 0){
            EVP_DigestSignUpdate(ctx, seed2, seed2_len);
        }
        if(seed3 != NULL && seed3_len > 0){
            EVP_DigestSignUpdate(ctx, seed3, seed3_len);
        }
        if(seed4 != NULL && seed4_len > 0){
            EVP_DigestSignUpdate(ctx, seed4, seed4_len);
        }

        if(remaining > chunk){
            EVP_DigestSignFinal(ctx, out, &len);
            out += len;
            remaining -= len;

            /* Next A value */
            EVP_DigestSignFinal(ctx_tmp, A, &A_len);
        } else {
            EVP_DigestSignFinal(ctx, A, &A_len);
            memcpy(out, A, remaining);
            remaining -= remaining;
        }
    }

    EVP_PKEY_free(mac_key);

#if OPENSSL_VERSION_NUMBER >= 0x1010000eL
    EVP_MD_CTX_free(ctx);
    EVP_MD_CTX_free(ctx_tmp);
    EVP_MD_CTX_free(ctx_init);
#else
    EVP_MD_CTX_cleanup(ctx);
    EVP_MD_CTX_cleanup(ctx_tmp);
    EVP_MD_CTX_cleanup(ctx_init);
    free(ctx);
    free(ctx_tmp);
    free(ctx_init);
#endif

    return 1;
}


/*
 * Generate the keys for the super encryption layer, based on the slitheen ID
 */
int generate_super_keys(uint8_t *secret){

    super = calloc(1, sizeof(super_data));

    EVP_MD_CTX *mac_ctx;

    const EVP_MD *md = EVP_sha256();

    /* Generate Keys */
    uint8_t *hdr_key, *bdy_key;
    uint8_t *mac_secret;
    EVP_PKEY *mac_key;
    int32_t mac_len, key_len;

    key_len = EVP_CIPHER_key_length(EVP_aes_256_cbc());
    mac_len = EVP_MD_size(md);
    int32_t total_len = 2*key_len + mac_len;
    uint8_t *key_block = calloc(1, total_len);

    PRF(secret, SLITHEEN_SUPER_SECRET_SIZE,
            (uint8_t *) SLITHEEN_SUPER_CONST, SLITHEEN_SUPER_CONST_SIZE,
            NULL, 0,
            NULL, 0,
            NULL, 0,
            key_block, total_len);

#ifdef DEBUG
    int i;
    printf("secret: \n");
    for(i=0; i< SLITHEEN_SUPER_SECRET_SIZE; i++){
        printf("%02x ", secret[i]);
    }
    printf("\n");
    printf("keyblock: \n");
    for(i=0; i< total_len; i++){
        printf("%02x ", key_block[i]);
    }
    printf("\n");
#endif

    hdr_key = key_block;
    bdy_key = key_block + key_len;
    mac_secret = key_block + 2*key_len;


    /* Initialize MAC Context */
    mac_ctx = EVP_MD_CTX_create();

    EVP_DigestInit_ex(mac_ctx, md, NULL);
    mac_key = EVP_PKEY_new_mac_key(EVP_PKEY_HMAC, NULL, mac_secret, mac_len);
    EVP_DigestSignInit(mac_ctx, NULL, md, NULL, mac_key);

    super->header_key = malloc(key_len);
    super->body_key = malloc(key_len);
    memcpy(super->header_key, hdr_key, key_len);
    memcpy(super->body_key, bdy_key, key_len);
    super->body_mac_ctx = mac_ctx;

    //Free everything
    free(key_block);
    EVP_PKEY_free(mac_key);

    return 0;
}

int peek_header(uint8_t *data){

    EVP_CIPHER_CTX *hdr_ctx = NULL;

    int32_t out_len;
    uint8_t *p = data;
    int retval = 1;

    //decrypt header
#ifdef DEBUG_PEEK
    int i;
    printf("Encrypted header:\n");
    for(i=0; i< SLITHEEN_HEADER_LEN; i++){
        printf("%02x ", p[i]);
    }
    printf("\n");
#endif

    hdr_ctx = EVP_CIPHER_CTX_new();

    EVP_CipherInit_ex(hdr_ctx, EVP_aes_256_ecb(), NULL, super->header_key, NULL, 0);

    if(!EVP_CipherUpdate(hdr_ctx, p, &out_len, p, SLITHEEN_HEADER_LEN)){
        retval =  0;
        goto end;
    }

    struct slitheen_hdr *sl_hdr = (struct slitheen_hdr *) p;

    //Finally, check to see if the last few bytes are zeroed.
    if(sl_hdr->zeros){
        //zeros have been corrupted, decryption failed
        printf("Decryption failed!\n");
        retval = 0;
        goto end;
    }

#ifdef DEBUG_PARSE
    printf("Decrypted header (%d bytes):\n", SLITHEEN_HEADER_LEN);
    for(i=0; i< SLITHEEN_HEADER_LEN; i++){
        printf("%02x ", p[i]);
    }
    printf("\n");
    fflush(stdout);
#endif

    retval = 1;

end:
    if(hdr_ctx != NULL){
        EVP_CIPHER_CTX_cleanup(hdr_ctx);
        OPENSSL_free(hdr_ctx);
    }

    return retval;

}

int super_decrypt(uint8_t *data){

    EVP_CIPHER_CTX *bdy_ctx = NULL;
    EVP_CIPHER_CTX *hdr_ctx = NULL;

    uint8_t *p = data;
    int32_t out_len, len;
    uint8_t output[EVP_MAX_MD_SIZE];
    size_t mac_len;
    int i, retval = 1;

    //decrypt header
#ifdef DEBUG
    printf("Encrypted header:\n");
    for(i=0; i< SLITHEEN_HEADER_LEN; i++){
        printf("%02x ", p[i]);
    }
    printf("\n");
#endif

    hdr_ctx = EVP_CIPHER_CTX_new();

    EVP_CipherInit_ex(hdr_ctx, EVP_aes_256_ecb(), NULL, super->header_key, NULL, 0);

    if(!EVP_CipherUpdate(hdr_ctx, p, &out_len, p, SLITHEEN_HEADER_LEN)){
        printf("Decryption failed!");
        retval =  0;
        goto end;
    }

    struct slitheen_hdr *sl_hdr = (struct slitheen_hdr *) p;
    len = htons(sl_hdr->len);

    if(!sl_hdr->len){//there are no data to be decrypted
        retval =  1;
        goto end;
    }

    if(len %16){ //add padding to len
        len += 16 - len%16;
    }


#ifdef DEBUG
    printf("Decrypted header (%d bytes):\n", SLITHEEN_HEADER_LEN);
    for(i=0; i< SLITHEEN_HEADER_LEN; i++){
        printf("%02x ", p[i]);
    }
    printf("\n");
    fflush(stdout);
#endif

    p += SLITHEEN_HEADER_LEN;

    //initialize body cipher context with IV
    bdy_ctx = EVP_CIPHER_CTX_new();
    EVP_CipherInit_ex(bdy_ctx, EVP_aes_256_cbc(), NULL, super->body_key, p, 0);
    p+=16;

    //compute mac
    EVP_MD_CTX *mac_ctx;
#if OPENSSL_VERSION_NUMBER >= 0x1010000eL
    mac_ctx = EVP_MD_CTX_new();
#else
    mac_ctx = ecalloc(1, sizeof(EVP_MD_CTX));
    EVP_MD_CTX_init(mac_ctx);
#endif
    EVP_MD_CTX_copy_ex(mac_ctx, super->body_mac_ctx);

    EVP_DigestSignUpdate(mac_ctx, p, len);

    EVP_DigestSignFinal(mac_ctx, output, &mac_len);

#if OPENSSL_VERSION_NUMBER >= 0x1010000eL
    EVP_MD_CTX_free(mac_ctx);
#else
    EVP_MD_CTX_cleanup(mac_ctx);
    free(mac_ctx);
#endif

#ifdef DEBUG_PARSE
    printf("Received mac:\n");
    for(i=0; i< 16; i++){
        printf("%02x ", p[len+i]);
    }
    printf("\n");
    fflush(stdout);
#endif

#ifdef DEBUG_PARSE
    printf("Computed mac:\n");
    for(i=0; i< 16; i++){
        printf("%02x ", output[i]);
    }
    printf("\n");
    fflush(stdout);
#endif

    if(memcmp(p+len, output, 16)){
        printf("MAC verification failed\n");
        retval =  0;
        goto end;
    }

    //decrypt body
#ifdef DEBUG_PARSE
    printf("Encrypted data (%d bytes):\n", len);
    for(i=0; i< len; i++){
        printf("%02x ", p[i]);
    }
    printf("\n");
#endif

    if(!EVP_CipherUpdate(bdy_ctx, p, &out_len, p, len)){
        printf("Decryption failed!");
        retval =  0;
        goto end;
    }

#ifdef DEBUG
    printf("Decrypted data (%d bytes):\n", out_len);
    for(i=0; i< out_len; i++){
        printf("%02x ", p[i]);
    }
    printf("\n");
    fflush(stdout);
#endif

    p += out_len;

    retval = 1;

end:
    if(hdr_ctx != NULL){
        EVP_CIPHER_CTX_cleanup(hdr_ctx);
        OPENSSL_free(hdr_ctx);
    }
    if(bdy_ctx != NULL){
        EVP_CIPHER_CTX_cleanup(bdy_ctx);
        OPENSSL_free(bdy_ctx);
    }

    return retval;

}