/*
 * Copyright 2009 Sun Microsystems, Inc.  All rights reserved.
 * Use is subject to license terms.
 */

/*
 * Copyright (c) 2003-2004, Jouni Malinen <jkmaline@cc.hut.fi>
 * Sun elects to license this software under the BSD license.
 * See README for more details.
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <strings.h>
#include <sys/types.h>

#include <openssl/aes.h>
#include <openssl/hmac.h>
#include <openssl/rc4.h>

#include "wpa_enc.h"

/*
 * @kek: key encryption key (KEK)
 * @n: length of the wrapped key in 64-bit units; e.g., 2 = 128-bit = 16 bytes
 * @plain: plaintext key to be wrapped, n * 64 bit
 * @cipher: wrapped key, (n + 1) * 64 bit
 */
void
aes_wrap(uint8_t *kek, int n, uint8_t *plain, uint8_t *cipher)
{
	uint8_t *a, *r, b[16];
	int i, j;
	AES_KEY key;

	a = cipher;
	r = cipher + 8;

	/* 1) Initialize variables. */
	(void) memset(a, 0xa6, 8);
	(void) memcpy(r, plain, 8 * n);

	(void) AES_set_encrypt_key(kek, 128, &key);

	/*
	 * 2) Calculate intermediate values.
	 * For j = 0 to 5
	 * 	For i=1 to n
	 * 		B = AES(K, A | R[i])
	 * 		A = MSB(64, B) ^ t where t = (n*j)+i
	 * 		R[i] = LSB(64, B)
	 */
	for (j = 0; j <= 5; j++) {
		r = cipher + 8;
		for (i = 1; i <= n; i++) {
			(void) memcpy(b, a, 8);
			(void) memcpy(b + 8, r, 8);
			AES_encrypt(b, b, &key);
			(void) memcpy(a, b, 8);
			a[7] ^= n * j + i;
			(void) memcpy(r, b + 8, 8);
			r += 8;
		}
	}

	/*
	 * 3) Output the results.
	 *
	 * These are already in @cipher due to the location of temporary
	 * variables.
	 */
}

/*
 * @kek: key encryption key (KEK)
 * @n: length of the wrapped key in 64-bit units; e.g., 2 = 128-bit = 16 bytes
 * @cipher: wrapped key to be unwrapped, (n + 1) * 64 bit
 * @plain: plaintext key, n * 64 bit
 */
int
aes_unwrap(uint8_t *kek, int n, uint8_t *cipher, uint8_t *plain)
{
	uint8_t a[8], *r, b[16];
	int i, j;
	AES_KEY key;

	/* 1) Initialize variables. */
	(void) memcpy(a, cipher, 8);
	r = plain;
	(void) memcpy(r, cipher + 8, 8 * n);

	(void) AES_set_decrypt_key(kek, 128, &key);

	/*
	 * 2) Compute intermediate values.
	 * For j = 5 to 0
	 * 	For i = n to 1
	 * 		B = AES-1(K, (A ^ t) | R[i]) where t = n*j+i
	 * 		A = MSB(64, B)
	 * 		R[i] = LSB(64, B)
	 */
	for (j = 5; j >= 0; j--) {
		r = plain + (n - 1) * 8;
		for (i = n; i >= 1; i--) {
			(void) memcpy(b, a, 8);
			b[7] ^= n * j + i;

			(void) memcpy(b + 8, r, 8);
			AES_decrypt(b, b, &key);
			(void) memcpy(a, b, 8);
			(void) memcpy(r, b + 8, 8);
			r -= 8;
		}
	}

	/*
	 * 3) Output results.
	 *
	 * These are already in @plain due to the location of temporary
	 * variables. Just verify that the IV matches with the expected value.
	 */
	for (i = 0; i < 8; i++) {
		if (a[i] != 0xa6) {
			return (-1);
		}
	}

	return (0);
}

/* RFC 2104 */
void
hmac_sha1(unsigned char *key, unsigned int key_len,
    unsigned char *data, unsigned int data_len, unsigned char *mac)
{
	unsigned int mac_len = 0;
	(void) HMAC(EVP_sha1(), key, key_len, data, data_len, mac, &mac_len);
}


void
hmac_sha1_vector(unsigned char *key, unsigned int key_len, size_t num_elem,
    unsigned char *addr[], unsigned int *len, unsigned char *mac)
{
	unsigned char *buf, *ptr;
	int i, buf_len;

	buf_len = 0;
	for (i = 0; i < num_elem; i ++)
		buf_len += len[i];

	buf = malloc(buf_len);
	ptr = buf;

	for (i = 0; i < num_elem; i ++) {
		(void) memcpy(ptr, addr[i], len[i]);
		ptr += len[i];
	}

	hmac_sha1(key, key_len, buf, buf_len, mac);

	free(buf);
}


void
sha1_prf(unsigned char *key, unsigned int key_len,
    char *label, unsigned char *data, unsigned int data_len,
    unsigned char *buf, size_t buf_len)
{
	uint8_t zero = 0, counter = 0;
	size_t pos, plen;
	uint8_t hash[SHA1_MAC_LEN];
	size_t label_len = strlen(label);

	unsigned char *addr[4];
	unsigned int len[4];

	addr[0] = (uint8_t *)label;
	len[0] = label_len;
	addr[1] = &zero;
	len[1] = 1;
	addr[2] = data;
	len[2] = data_len;
	addr[3] = &counter;
	len[3] = 1;

	pos = 0;
	while (pos < buf_len) {
		plen = buf_len - pos;
		if (plen >= SHA1_MAC_LEN) {
			hmac_sha1_vector(key, key_len, 4, addr, len, &buf[pos]);
			pos += SHA1_MAC_LEN;
		} else {
			hmac_sha1_vector(key, key_len, 4, addr, len, hash);
			(void) memcpy(&buf[pos], hash, plen);
			break;
		}
		counter++;
	}
}

void
pbkdf2_sha1(char *passphrase, char *ssid, size_t ssid_len, int iterations,
    unsigned char *buf, size_t buflen)
{
	(void) PKCS5_PBKDF2_HMAC_SHA1(passphrase, -1, (unsigned char *)ssid,
	    ssid_len, iterations, buflen, buf);
}

void
rc4_skip(uint8_t *key, size_t keylen, size_t skip,
    uint8_t *data, size_t data_len)
{
	uint8_t *buf;
	size_t buf_len;

	buf_len = skip + data_len;
	buf = malloc(buf_len);

	bzero(buf, buf_len);
	bcopy(data, buf + skip, data_len);

	rc4(buf, buf_len, key, keylen);

	bcopy(buf + skip, data, data_len);
	free(buf);
}

void
rc4(uint8_t *buf, size_t len, uint8_t *key, size_t key_len)
{
	RC4_KEY k;

	RC4_set_key(&k, key_len, key);
	RC4(&k, len, buf, buf);
}

void
hmac_md5_vector(uint8_t *key, size_t key_len, size_t num_elem,
    uint8_t *addr[], size_t *len, uint8_t *mac)
{
	unsigned char *buf, *ptr;
	int i, buf_len;

	buf_len = 0;
	for (i = 0; i < num_elem; i ++)
		buf_len += len[i];

	buf = malloc(buf_len);
	ptr = buf;

	for (i = 0; i < num_elem; i ++) {
		(void) memcpy(ptr, addr[i], len[i]);
		ptr += len[i];
	}

	hmac_md5(key, key_len, buf, buf_len, mac);
	free(buf);
}

/* RFC 2104 */
void
hmac_md5(uint8_t *key, size_t key_len, uint8_t *data,
    size_t data_len, uint8_t *mac)
{
	unsigned int mac_len = 0;
	(void) HMAC(EVP_md5(), key, key_len, data, data_len, mac, &mac_len);
}