#include <linux/init.h>				// Macros used to mark up functions e.g. __init __exit
#include <linux/module.h>			// Core header for loading LKMs into the kernel
#include <linux/device.h>			// Header to support the kernel Driver Model
#include <linux/kernel.h>			// Contains types, macros, functions for the kernel
#include <linux/fs.h>				// Header for the Linux file system support
#include <asm/uaccess.h>			// Required for the copy to user function

#include <crypto/skcipher.h>
#include <linux/scatterlist.h>
#include <crypto/hash.h>
#include <crypto/sha.h>

#define DEVICE_NAME "crypto"		///< The device will appear at /dev/crypto using this value
#define CLASS_NAME  "cipher"		///< The device class -- this is a character device driver

MODULE_LICENSE("GPL");
MODULE_AUTHOR("Marcelo Aquino, Jefferson Aparecido");
MODULE_DESCRIPTION("A crypto Linux char driver");  ///< The description -- see modinfo
MODULE_VERSION("0.1");            // A version number to inform users

struct tcrypt_result {
    struct completion completion;
    int res;
};

/* tie all data structures together */
struct skcipher_def {
    struct crypto_skcipher *tfm;
    struct skcipher_request *req;
    struct tcrypt_result result;
};

// Module parameters
static char *key = NULL;
module_param(key, charp, 0000);
MODULE_PARM_DESC(key, "Chave simetrica para cifra e decifrar os dados");

static int    majorNumber;                  ///< Stores the device number -- determined automatically
static char   buffer[1024] = {0};           ///< Memory for the string that is passed from userspace
static short  buffer_size;              ///< Used to remember the size of the string stored
static int    numberOpens = 0;              ///< Counts the number of times the device is opened
static struct class*  ebbcharClass  = NULL; ///< The device-driver class struct pointer
static struct device* ebbcharDevice = NULL; ///< The device-driver device struct pointer

// The prototype functions for the character driver -- must come before the struct definition
static int     dev_open(struct inode *, struct file *);
static int     dev_release(struct inode *, struct file *);
static ssize_t dev_read(struct file *, char *, size_t, loff_t *);
static ssize_t dev_write(struct file *, const char *, size_t, loff_t *);

// Crypto
static struct skcipher_def sk;
static struct crypto_shash *hashalg;

/** @brief Devices are represented as file structure in the kernel. The file_operations structure from
 *  /linux/fs.h lists the callback functions that you wish to associated with your file operations
 *  using a C99 syntax structure. char devices usually implement open, read, write and release calls
 */
static struct file_operations fops =
{
	.open = dev_open,
	.read = dev_read,
	.write = dev_write,
	.release = dev_release,
};

struct sdesc {
	struct shash_desc shash;
	char ctx[];
};

/* Callback function */
static void test_skcipher_cb(struct crypto_async_request *req, int rc)
{
    struct tcrypt_result *result = req->data;

    if (rc == -EINPROGRESS) {
#ifdef CRYPTO_DEBUG
		printk(KERN_INFO "crypto: failed\n");
#endif
        return;
	}

    result->res = rc;
    complete(&result->completion);

    pr_info("Encryption finished successfully\n");
}

/** @brief The LKM initialization function
 *  The static keyword restricts the visibility of the function to within this C file. The __init
 *  macro means that for a built-in driver (not a LKM) the function is only used at initialization
 *  time and that it can be discarded and its memory freed up after that point.
 *  @return returns 0 if successful
 */
static int __init crypto_init(void)
{
	int ret;

	printk(KERN_INFO "crypto: Initializing the crypto device\n");

	if (key == NULL) {
		printk(KERN_ERR "crypto: You must supply the AES key to the module.\n");
		return -1;
	} else if (strlen(key) != 32) {
		printk(KERN_ERR "crypto: You must supply an AES with size of 32.\n");
		return -1;
	} else {
		printk(KERN_INFO "crypto: AES key set to: %s\n", key);
	}

	sk.tfm = crypto_alloc_skcipher("ecb(aes)", 0, 0);
	if (IS_ERR(sk.tfm)) {
		printk(KERN_ERR "crypto: could not allocate skcipher handle\n");
		return PTR_ERR(sk.tfm);
	}

	sk.req = skcipher_request_alloc(sk.tfm, GFP_KERNEL);
    if (!sk.req) {
        printk(KERN_ERR "crypto: could not allocate skcipher request\n");
        return -ENOMEM;
    }

	skcipher_request_set_callback(sk.req, CRYPTO_TFM_REQ_MAY_BACKLOG,
		test_skcipher_cb, &sk.result);

	if (crypto_skcipher_setkey(sk.tfm, key, 32)) {
		printk(KERN_ERR "crypto: aes key could not be set\n");
		crypto_free_skcipher(sk.tfm);
		return -EAGAIN;
	}

	hashalg = crypto_alloc_shash("sha1", 0, CRYPTO_ALG_ASYNC);
	if (IS_ERR(hashalg)) {
		pr_info("crypto: could not allocate crypto sha1\n");
		ret = PTR_ERR(hashalg);
		goto hashalg_fail;
	}

	// Try to dynamically allocate a major number for the device -- more difficult but worth it
	majorNumber = register_chrdev(0, DEVICE_NAME, &fops);
	if (majorNumber<0){
		printk(KERN_ALERT "crypto failed to register a major number\n");
		ret = majorNumber;
		goto device_register_fail;
	}
	printk(KERN_INFO "crypto: registered correctly with major number %d\n", majorNumber);

	// Register the device class
	ebbcharClass = class_create(THIS_MODULE, CLASS_NAME);
	if (IS_ERR(ebbcharClass)){                // Check for error and clean up if there is
		unregister_chrdev(majorNumber, DEVICE_NAME);
		printk(KERN_ALERT "Failed to register device class\n");
		return PTR_ERR(ebbcharClass);          // Correct way to return an error on a pointer
	}
	printk(KERN_INFO "crypto: device class registered correctly\n");

	// Register the device driver
	ebbcharDevice = device_create(ebbcharClass, NULL, MKDEV(majorNumber, 0), NULL, DEVICE_NAME);
	if (IS_ERR(ebbcharDevice)){               // Clean up if there is an error
		class_destroy(ebbcharClass);           // Repeated code but the alternative is goto statements
		unregister_chrdev(majorNumber, DEVICE_NAME);
		printk(KERN_ALERT "Failed to create the device\n");
		return PTR_ERR(ebbcharDevice);
	}
	printk(KERN_INFO "crypto: device class created correctly\n"); // Made it! device was initialized

	return 0;

	hashalg_fail:
		crypto_free_skcipher(sk.tfm);
		return ret;
	device_register_fail:
		crypto_free_shash(hashalg);
		return ret;
}

/** @brief The LKM cleanup function
 *  Similar to the initialization function, it is static. The __exit macro notifies that if this
 *  code is used for a built-in driver (not a LKM) that this function is not required.
 */
static void __exit crypto_exit(void)
{
	if (sk.tfm != NULL) {
		crypto_free_skcipher(sk.tfm);
	}
	if (sk.req != NULL) {
		skcipher_request_free(sk.req);
	}

	device_destroy(ebbcharClass, MKDEV(majorNumber, 0));     // remove the device
	class_unregister(ebbcharClass);                          // unregister the device class
	class_destroy(ebbcharClass);                             // remove the device class
	unregister_chrdev(majorNumber, DEVICE_NAME);             // unregister the major number
	printk(KERN_INFO "crypto: Goodbye from the crypto module\n");
}

/** @brief The device open function that is called each time the device is opened
 *  This will only increment the numberOpens counter in this case.
 *  @param inodep A pointer to an inode object (defined in linux/fs.h)
 *  @param filep A pointer to a file object (defined in linux/fs.h)
 */
static int dev_open(struct inode *inodep, struct file *filep)
{
	numberOpens++;
	printk(KERN_INFO "crypto: Device has been opened %d time(s)\n", numberOpens);

	return 0;
}

/** @brief This function is called whenever device is being read from user space i.e. data is
 *  being sent from the device to the user. In this case is uses the copy_to_user() function to
 *  send the buffer string to the user and captures any errors.
 *  @param filep A pointer to a file object (defined in linux/fs.h)
 *  @param buffer The pointer to the buffer to which this function writes the data
 *  @param len The length of the b
 *  @param offset The offset if required
 */
static ssize_t dev_read(struct file *filep, char *ubuf, size_t len, loff_t *offset)
{
	int error_count = 0;
	// copy_to_user has the format ( * to, *from, size) and returns 0 on success
	error_count = copy_to_user(ubuf, buffer, buffer_size);

	if (error_count==0) {            // if true then have success
#ifdef CRYPTO_DEBUG
		printk(KERN_INFO "crypto: Sent %d characters to the user\n", buffer_size);
#endif
		return buffer_size;
	} else {
		printk(KERN_INFO "crypto: Failed to send %d characters to the user\n", error_count);
		return -EFAULT;              // Failed -- return a bad address buffer (i.e. -14)
	}
}

/** @brief Encrypt or decrypt a string
 *  @param op Operation, 0 being encrypt and 1 decrypt
 *  @param src The source string
 *  @param dst The destination string
 *  @param len The length of the source string
 */
static int crypto_func(uint8_t op, const char *src, char *dst, int len)
{
	char source[16];
	char derived[16];
	int i, rc, blocks, bytes = 0;
#ifdef CRYPTO_DEBUG
	int j;
#endif
	struct scatterlist src_sg, dst_sg;

	if (len == 0 || len % 16) {
#ifdef CRYPTO_DEBUG
		printk(KERN_INFO "crypto: invalid string size\n");
#endif
		return 0;
	}

	blocks = len / 16;

	sg_init_one(&src_sg, source, 16);
	sg_init_one(&dst_sg, derived, 16);
	skcipher_request_set_crypt(sk.req, &src_sg, &dst_sg, 16, NULL);

	for (i = 0; i < blocks; ++i) {
		(void)memcpy((void *)source, (void *)&src[bytes], 16);

		if (op == 0) {
#ifdef CRYPTO_DEBUG
			printk(KERN_INFO "crypto: encrypting block: %d\n", i);
			printk(KERN_INFO "crypto: source hex: ");
			for (j = 0; j < 16; ++j) {
				printk(KERN_CONT "0x%02X ", source[j] & 0xff);
			}
#endif

			rc = crypto_skcipher_encrypt(sk.req);
		} else {
#ifdef CRYPTO_DEBUG
			printk(KERN_INFO "crypto: decrypting block: %d\n", i);
			printk(KERN_INFO "crypto: source hex: ");
			for (j = 0; j < 16; ++j) {
				printk(KERN_CONT "0x%02X ", source[j] & 0xff);
			}
#endif
			rc = crypto_skcipher_decrypt(sk.req);
		}

		if (rc == -EINPROGRESS || rc == -EBUSY) {
			wait_for_completion(&sk.result.completion);
			rc = sk.result.res;
		}

#ifdef CRYPTO_DEBUG
		printk(KERN_INFO "crypto: derived hex: ");
		for (j = 0; j < 16; ++j) {
			printk(KERN_CONT "0x%02X ", derived[j] & 0xff);
		}
#endif
		(void)memcpy((void *)&dst[bytes], (void *)derived, 16);
		bytes += 16;
	}

	return bytes;
}

uint8_t h2i(char c)
{
	uint8_t i = 0;
	if (c <= '9') {
		i += c - '0';
	} else if (c >= 'a') {
		i += c - 'a' + 10;
	} else {
		i += c - 'A' + 10;
	}
	return i;
}

char i2h(uint8_t i)
{
	uint8_t k = i & 0x0F;
	if (k <= 9) {
		return '0' + k;
	} else {
		return 'a' + k - 10;
	}
}

static struct sdesc *init_sdesc(struct crypto_shash *alg)
{
	struct sdesc *sdesc;
	int size;

	size = sizeof(struct shash_desc) + crypto_shash_descsize(alg);
	sdesc = kmalloc(size, GFP_KERNEL);
	if (!sdesc) {
		return ERR_PTR(-ENOMEM);
	}
	sdesc->shash.tfm = alg;
	sdesc->shash.flags = 0x0;
	return sdesc;
}

static int TSS_sha1(const unsigned char *data, unsigned int datalen,
		    unsigned char *digest)
{
	struct sdesc *sdesc;
	int ret;

	sdesc = init_sdesc(hashalg);
	if (IS_ERR(sdesc)) {
		pr_info("trusted_key: can't alloc sha1\n");
		return PTR_ERR(sdesc);
	}

	ret = crypto_shash_digest(&sdesc->shash, data, datalen, digest);
	kfree(sdesc);
	return ret;
}

/** @brief This function is called whenever the device is being written to from user space i.e.
 *  data is sent to the device from the user. The data is copied to the buffer[] array in this
 *  LKM using the sprintf() function along with the length of the string.
 *  @param filep A pointer to a file object
 *  @param buffer The buffer to that contains the string to write to the device
 *  @param len The length of the array of data that is being passed in the const char buffer
 *  @param offset The offset if required
 */
static ssize_t dev_write(struct file *filep, const char *ubuf, size_t len, loff_t *offset)
{
	int i, rc = 0;
	uint8_t padding, val;
	char *str;
	const uint8_t zero = 0;
	unsigned char digest[SHA1_DIGEST_SIZE] = {0};

	if (len < 2) {
		return -1;
	}
#ifdef CRYPTO_DEBUG
	printk(KERN_INFO "crypto: write: ubuf: '%.*s' - len: %zu\n", (int)len, ubuf, len);
#endif
	len = len - 2;
	if (!len) {
		return 0;
	}

	if (ubuf[0] == 'c' && ubuf[1] == ' ') {
		buffer_size = 0;

		(void)memcpy((void *)buffer, (void *)&ubuf[2], len);

		padding = (16 - len % 16) % 16;
		if (padding) {
			for (i = len; i < len + padding; ++i) {
				(void)memcpy((void *)&buffer[i], &zero, 1);
			}
			len = len + padding;
		}
		rc = crypto_func(0, buffer, buffer, len);
		if (rc > 0) {
			str = &buffer[(len - 1) * 2];
			for (i = len - 1; i >= 0; --i) {
				val = buffer[i];
				buffer[i * 2] = i2h(val >> 4);
				buffer[(i * 2) + 1] = i2h(val);
			}
			buffer_size = len * 2;
#ifdef CRYPTO_DEBUG
			printk(KERN_INFO "crypto: buffer - size: %d - value: %.*s", buffer_size, buffer_size, buffer);
#endif
		}
	} else if (ubuf[0] == 'd' && ubuf[1] == ' ') {
		if (len % (16*2)) {
			return -1;
		}
		str = buffer;
		for (i = 2; i <= len; str++) {
			val = h2i(ubuf[i++]) << 4;
			val += h2i(ubuf[i++]);
			*str = val;
		}
		rc = crypto_func(1, buffer, buffer, len/2);
		if (rc > 0) {
			buffer_size = rc;
			printk(KERN_INFO "crypto: buffer - size: %d - value: %.*s", buffer_size, buffer_size, buffer);
		} else {
			buffer_size = 0;
		}
	} else if (ubuf[0] == 'h' && ubuf[1] == ' ') {
		buffer_size = 0;
		rc = TSS_sha1(&ubuf[2], len, digest);
		if (rc < 0) {
#ifdef CRYPTO_DEBUG
			printk(KERN_INFO "crypto: sha1 failed\n");
#endif
		} else {
			for (i = 0; i < SHA1_DIGEST_SIZE; ++i) {
				val = digest[i];
				buffer[i * 2] = i2h(val >> 4);
				buffer[(i * 2) + 1] = i2h(val);
			}
			buffer_size = SHA1_DIGEST_SIZE * 2;
		}
	}

	return rc;
}

/** @brief The device release function that is called whenever the device is closed/released by
 *  the userspace program
 *  @param inodep A pointer to an inode object (defined in linux/fs.h)
 *  @param filep A pointer to a file object (defined in linux/fs.h)
 */
static int dev_release(struct inode *inodep, struct file *filep)
{
#ifdef CRYPTO_DEBUG
	printk(KERN_INFO "crypto: Device successfully closed\n");
#endif

	return 0;
}

/** @brief A module must use the module_init() module_exit() macros from linux/init.h, which
 *  identify the initialization function at insertion time and the cleanup function (as
 *  listed above)
 */
module_init(crypto_init);
module_exit(crypto_exit);