/*
version 20080913
D. J. Bernstein
Public domain.
*/

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <openssl/rand.h>
#include <openssl/evp.h>
#include <crypto_scrypt.h>
#include "crypto_core_hsalsa20.h"
#include "crypto_stream_salsa20.h"
#include "crypto_xsalsa20.h"

extern int geturandom_bytes(uchar_t *rbytes, int nbytes);

static const unsigned char sigma[16] = "expand 32-byte k";
static const unsigned char tau[16] = "expand 16-byte k";

static int
crypto_salsa20(unsigned char *c, const unsigned char *m, unsigned long long mlen,
  const unsigned char *n, const unsigned char *k, int klen)
{
	unsigned char subkey[32];
	if (klen < XSALSA20_CRYPTO_KEYBYTES)
		crypto_core_hsalsa20(subkey,n,k,tau);
	else
		crypto_core_hsalsa20(subkey,n,k,sigma);
	return crypto_stream_salsa20_xor(c,m,mlen,n + 16,subkey);
}

int
salsa20_init(salsa20_ctx_t *ctx, uchar_t *salt, int saltlen, uchar_t *pwd, int pwd_len,
	 uchar_t *nonce, int enc)
{
	struct timespec tp;
	uint64_t tv;
	uchar_t num[25];
	uchar_t IV[32];
	uchar_t *key = ctx->pkey;

#ifndef	_USE_PBK
	int logN;
	uint32_t r, p;
	uint64_t N;

	if (XSALSA20_CRYPTO_NONCEBYTES % 8) {
		fprintf(stderr, "XSALSA20_CRYPTO_NONCEBYTES is not a multiple of 8!\n");
		return (-1);
	}
	pickparams(&logN, &r, &p);
	N = (uint64_t)(1) << logN;
	if (crypto_scrypt(pwd, pwd_len, salt, saltlen, N, r, p, key, ctx->keylen)) {
		fprintf(stderr, "Scrypt failed\n");
		return (-1);
	}
#else
	rv = PKCS5_PBKDF2_HMAC(pwd, pwd_len, salt, saltlen, PBE_ROUNDS, EVP_sha256(),
			       ctx->keylen, key);
	if (rv != ctx->keylen) {
		fprintf(stderr, "Key size is %d bytes - should be %d bits\n", i, ctx->keylen);
		return (-1);
	}
#endif

	/*
	 * Copy the key. XSalsa20 core cipher always uses a 256-bit key. If we are using a
	 * 128-bit key then the key value is repeated twice to form a 256-bit value.
	 * This approach is based on the Salsa20 code submitted to eSTREAM. See the function
	 * ECRYPT_keysetup() in the Salsa20 submission:
	 * http://www.ecrypt.eu.org/stream/svn/viewcvs.cgi/ecrypt/trunk/submissions/salsa20/full/ref/salsa20.c?rev=161&view=auto
	 * 
	 * The input values corresponding to a 256-bit key contain repeated values if key
	 * length is 128-bit.
	 */
	memcpy(ctx->key, key, ctx->keylen);
	if (ctx->keylen < XSALSA20_CRYPTO_KEYBYTES) {
		uchar_t *k;
		k = ctx->key + ctx->keylen;
		memcpy(k, key, XSALSA20_CRYPTO_KEYBYTES - ctx->keylen);
	}

	if (enc) {
		int i;
		uint64_t *n, *n1;

		// Derive 192-bit nonce
		if (RAND_status() != 1 || RAND_bytes(IV, XSALSA20_CRYPTO_NONCEBYTES) != 1) {
			if (geturandom_bytes(IV, XSALSA20_CRYPTO_NONCEBYTES) != 0) {
				if (clock_gettime(CLOCK_MONOTONIC, &tp) == -1) {
					time((time_t *)&tv);
				} else {
					tv = tp.tv_sec * 1000UL + tp.tv_nsec;
				}
				sprintf((char *)num, "%" PRIu64, tv);
				PKCS5_PBKDF2_HMAC((const char *)num, strlen((char *)num), salt,
						saltlen, PBE_ROUNDS, EVP_sha256(), 32, IV);
			}
		}
		n = (uint64_t *)IV;
		n1 = (uint64_t *)(ctx->nonce);
		for (i = 0; i < XSALSA20_CRYPTO_NONCEBYTES/8; i++) {
			*n1 = ntohll(*n);
			n++;
			n1++;
		}

		// Nullify stack components
		memset(num, 0, 25);
		memset(IV, 0, 32);
		memset(&tp, 0, sizeof (tp));
		tv = 0;
	} else {
		memcpy(ctx->nonce, nonce, XSALSA20_CRYPTO_NONCEBYTES);
		memset(nonce, 0, XSALSA20_CRYPTO_NONCEBYTES);
	}
	return (0);
}

int
salsa20_encrypt(salsa20_ctx_t *ctx, uchar_t *plaintext, uchar_t *ciphertext, uint64_t len, uint64_t id)
{
	uchar_t nonce[XSALSA20_CRYPTO_NONCEBYTES];
	int i;
	uint64_t *n, carry;

	for (i = 0; i < XSALSA20_CRYPTO_NONCEBYTES; i++) nonce[i] = ctx->nonce[i];
	carry = id;
	n = (uint64_t *)nonce;
	for (i = 0; i < XSALSA20_CRYPTO_NONCEBYTES/8; i++) {
		if (UINT64_MAX - *n < carry) {
			carry = carry - (UINT64_MAX - *n);
			*n = 0;
		} else {
			*n += carry;
			carry = 0;
			break;
		}
		n++;
	}
	if (carry) {
		n = (uint64_t *)nonce;
		*n += carry;
	}

	return (crypto_salsa20(ciphertext, plaintext, len, nonce, ctx->key, ctx->keylen));
}

int
salsa20_decrypt(salsa20_ctx_t *ctx, uchar_t *ciphertext, uchar_t *plaintext, uint64_t len, uint64_t id)
{
	uchar_t nonce[XSALSA20_CRYPTO_NONCEBYTES];
	int i;
	uint64_t *n, carry;

	for (i = 0; i < XSALSA20_CRYPTO_NONCEBYTES; i++) nonce[i] = ctx->nonce[i];
	carry = id;
	n = (uint64_t *)nonce;
	for (i = 0; i < XSALSA20_CRYPTO_NONCEBYTES/8; i++) {
		if (UINT64_MAX - *n < carry) {
			carry = carry - (UINT64_MAX - *n);
			*n = 0;
		} else {
			*n += carry;
			carry = 0;
			break;
		}
		n++;
	}
	if (carry) {
		n = (uint64_t *)nonce;
		*n += carry;
	}

	return (crypto_salsa20(plaintext, ciphertext, len, nonce, ctx->key, ctx->keylen));
}

uchar_t *
salsa20_nonce(salsa20_ctx_t *ctx)
{
	return (ctx->nonce);
}

void
salsa20_clean_pkey(salsa20_ctx_t *ctx)
{
	memset(ctx->pkey, 0, ctx->keylen);
}

void
salsa20_cleanup(salsa20_ctx_t *ctx)
{
	memset((void *)(&ctx->key), 0, sizeof (ctx->key));
	memset(ctx->nonce, 0, XSALSA20_CRYPTO_NONCEBYTES);
	free(ctx);
}