#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "pbkdf2_hmac.h"
#include "sha.h"


#ifndef min
#define min( a, b ) ( ((a) < (b)) ? (a) : (b) )
#endif


void sha1_hmac_starts(sha1_context * ctx, const unsigned char *key, int keylen)
{
    int i;
    unsigned char sum[20];

    if (keylen > 64) {
        SHA_Simple(key, keylen, sum);
        keylen = 20;
        key = sum;
    }

    memset(ctx->ipad, 0x36, KEY_IOPAD_SIZE);
    memset(ctx->opad, 0x5C, KEY_IOPAD_SIZE);

    for (i = 0; i < keylen; i++) {
        ctx->ipad[i] = (unsigned char)(ctx->ipad[i] ^ key[i]);
        ctx->opad[i] = (unsigned char)(ctx->opad[i] ^ key[i]);
    }

    SHA_Init(&ctx->ctx);
    SHA_Bytes(&ctx->ctx, ctx->ipad, KEY_IOPAD_SIZE);

}

/*
* SHA-1 HMAC process buffer
*/
void sha1_hmac_update(sha1_context * ctx, const unsigned char *input, int ilen)
{
    SHA_Bytes(&ctx->ctx, input, ilen);
}

/*
* SHA-1 HMAC final digest
*/
void sha1_hmac_finish(sha1_context * ctx, unsigned char output[20])
{
    unsigned char tmpbuf[20];

    SHA_Final(&ctx->ctx, tmpbuf);
    SHA_Init(&ctx->ctx);
    SHA_Bytes(&ctx->ctx, ctx->opad, 64);
    SHA_Bytes(&ctx->ctx, tmpbuf, 20);
    SHA_Final(&ctx->ctx, output);

}


void PKCS5_PBKDF2_HMAC(const unsigned char *password, size_t plen,
    const unsigned char *salt, size_t slen,
    const unsigned long iteration_count, const unsigned long key_length,
    unsigned char *output)
{
    sha1_context ctx;
    SHA_Init(&ctx.ctx);

	// Size of the generated digest
	unsigned char md_size = 20;
	unsigned char md1[20];
	unsigned char work[20];

	unsigned long counter = 1;
	unsigned long generated_key_length = 0;
	while (generated_key_length < key_length) {
		// U1 ends up in md1 and work
		unsigned char c[4];
        PUT_32BIT_MSB_FIRST(c, counter);
        sha1_hmac_starts(&ctx, password, plen);
        sha1_hmac_update(&ctx, salt, slen);
        sha1_hmac_update(&ctx, c, 4);
        sha1_hmac_finish(&ctx, md1);
		memcpy(work, md1, md_size);

		unsigned long ic = 1;
		for (ic = 1; ic < iteration_count; ic++) {
			// U2 ends up in md1
            sha1_hmac_starts(&ctx, password, plen);
            sha1_hmac_update(&ctx, md1, md_size);
            sha1_hmac_finish(&ctx, md1);
			// U1 xor U2
			unsigned long i = 0;
			for (i = 0; i < md_size; i++) {
				work[i] ^= md1[i];
			}
			// and so on until iteration_count
		}

		// Copy the generated bytes to the key
		unsigned long bytes_to_write =
		    min((key_length - generated_key_length), md_size);
		memcpy(output + generated_key_length, work, bytes_to_write);
		generated_key_length += bytes_to_write;
		++counter;
	}
}

