/*
 * 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;

}