#include <crypto/hash.h>
#include <linux/err.h>
#include <linux/module.h>
#include <linux/scatterlist.h>
#include <linux/slab.h>
#include <linux/string.h>
#include <crypto/rng.h>
#include <crypto/authenc.h>
#include <linux/pagemap.h>
#include <linux/rtnetlink.h>

#define MAX_IVLEN		16
#define AUTHSIZE		20

int testmode = 0;

module_param(testmode, int, 0);

void hexdump(unsigned char *buf, int len) 
{
	while(--len >=0)
		printk("%02X", *buf++);
	printk("\n");
}


struct tcrypt_result {
	struct completion completion;
	int err;
};


static void tcrypt_complete(struct crypto_async_request *req, int err)
{
	struct tcrypt_result *res = req->data;

	if (err == -EINPROGRESS)
		return;

	res->err = err;
	complete(&res->completion);
}


static int testrun(void)
{
 	const char *algo = "authenc(hmac(sha1),cbc(aes))";

	struct crypto_aead *tfm;
	unsigned int i;
	int ret = -ENOMEM;
	unsigned char *ekey = "1234567890123456";
	unsigned char *akey = "123456789012";
	unsigned char *key, *pp;
	unsigned int ekeylen = 16;
	unsigned int akeylen = 12;
	unsigned int keylen = 0;
	struct aead_request *req;
	struct tcrypt_result result;
	unsigned int authsize = AUTHSIZE;
	char *assoc = "\x49\x5c\x50\x1f\x1d\x94\xcc\x81"
			  "\xba\xb7\xb6\x03\xaf\xa5\xc1\xa1"
			  "\xd8\x5c\x42\x68\xe0\x6c\xda\x89"
			  "\x05\xac\x56\xac\x1b\x2a\xd3\x86";
	int assoclen = 32;
	unsigned char iv[MAX_IVLEN];
	mm_segment_t old_fs = get_fs();
	struct file *fp;
	struct scatterlist src_sg, dst_sg;
	struct scatterlist asg;
	char *filename = "file1.txt";
	char *string = "ciao";
	struct crypto_authenc_key_param *param;
	struct rtattr *rta;
	char *input;
	int blksize = 0;
	int clen = 16;		/* len of data */

	printk("assoc: ");
	hexdump(assoc, 32);

	tfm = crypto_alloc_aead(algo, 0, 0);
	if (IS_ERR(tfm)) {
		printk(KERN_ERR "alg: aead: Failed to load transform for %s: "
		       "%ld\n", algo, PTR_ERR(tfm));
		goto out;
	}

	blksize = ALIGN(crypto_aead_blocksize(tfm), 4);
	clen = ALIGN(clen + 2, blksize);
	input = kzalloc(clen + AUTHSIZE, GFP_KERNEL);
	
	if(!testmode) {
		set_fs(KERNEL_DS);

		fp = filp_open(filename, O_RDONLY, 0);
		if(IS_ERR(fp)) {
			printk("Error opening file: %s\n", filename);
			set_fs(old_fs);
			goto out;
		}
		fp->f_op->read(fp, input, clen + AUTHSIZE, &fp->f_pos);

		filp_close(fp, NULL); 
		set_fs(old_fs);
	} else
		memcpy(input, string, strlen(string));	

	init_completion(&result.completion);
	
	req = aead_request_alloc(tfm, GFP_KERNEL);
	if (!req) {
		printk(KERN_ERR "alg: aead: Failed to allocate request for "
		       "%s\n", algo);
		goto out_tfm;
	}
	aead_request_set_callback(req, CRYPTO_TFM_REQ_MAY_BACKLOG,
				  tcrypt_complete, &result);

	memset(iv, 1, MAX_IVLEN);
	crypto_aead_clear_flags(tfm, ~0);
	crypto_aead_set_flags(tfm, CRYPTO_TFM_REQ_WEAK_KEY);
	
	keylen = akeylen + ekeylen + RTA_SPACE(sizeof(*param));
	key = kzalloc(keylen, GFP_KERNEL);
	if (!key)
		goto out_req;

	pp = key;
	rta = (void *)pp;
	rta->rta_type = CRYPTO_AUTHENC_KEYA_PARAM;
	rta->rta_len = RTA_LENGTH(sizeof(*param));
	param = RTA_DATA(rta);
	pp += RTA_SPACE(sizeof(*param));

	if(akeylen) {
		memcpy(pp, akey, akeylen);
		pp += akeylen;
	}

	ret = crypto_aead_setauthsize(tfm, authsize);
	if (ret) {
		printk(KERN_ERR "alg: aead: Failed to set "
		       "authsize to %u on test for %s\n",
		       authsize, algo);
		goto out_key;
	}
	
	param->enckeylen = cpu_to_be32(ekeylen);
	memcpy(pp, ekey, ekeylen);
	
	ret = crypto_aead_setkey(tfm, key, keylen);
	if (ret) {
		printk(KERN_ERR "alg: aead: Failed to set "
		       "the key for %s\n",
		       algo);
		goto out_key;
	}

	sg_init_one(&src_sg, input, clen + AUTHSIZE);
	sg_init_one(&dst_sg, input, clen + AUTHSIZE);
	sg_init_one(&asg, assoc, assoclen);

	aead_request_set_crypt(req, &src_sg, &dst_sg, clen +  (testmode ? 0 : authsize), iv);

	aead_request_set_assoc(req, &asg, assoclen);
	
	ret = testmode ?
		crypto_aead_encrypt(req) :
		crypto_aead_decrypt(req);

	switch (ret) {
		case 0:
			if(!testmode)
				printk("File OK\n");
			break;
		case -EINPROGRESS:	  
		case -EBUSY:
			ret = wait_for_completion_interruptible(
				&result.completion);
			if (!ret && !(ret = result.err)) {
				INIT_COMPLETION(result.completion);
				break;
			}
		case -EBADMSG:
			printk("File corrupted\n");
		default:
			printk(KERN_ERR "alg: aead: %s failed "
						"for %s: ret=%d\n", testmode ? "encrypt" : "decrypt", algo, -ret);
	}
	
	if(ret)
		goto out_key;
	printk("%s: %d\n", testmode ? "encrypt" : "decrypt", ret);
	
	if(testmode) {
		set_fs(KERNEL_DS);

		fp = filp_open(filename, O_WRONLY | O_CREAT | O_TRUNC, 0700);
		if(IS_ERR(fp)) {
			printk("Error opening file: %s\n", filename);
			set_fs(old_fs);
			goto out_key;
		}
		fp->f_op->write(fp, input, clen + AUTHSIZE, &fp->f_pos);
		hexdump(input + clen, AUTHSIZE);
		filp_close(fp, NULL); 
		set_fs(old_fs);
	} else {
		printk("Decrypted: %s\n", input);
	}
	kfree(input);
out_key:
	kfree(key);
out_req:
	aead_request_free(req);
out_tfm:
	crypto_free_aead(tfm);
out:

	return ret;
}

int __init testaead_start(void) {
	int ret;
	printk("--------- testing authenc: %s ---------------\n", testmode ? "encrypt" : "decrypt");
	ret = testrun();
	printk("Result: %d\n", ret);
	printk("--------------------------------------------------\n");

	return -1;
}


void __exit testaead_stop(void) {

}

module_init(testaead_start);
module_exit(testaead_stop);
MODULE_LICENSE("GPL");