/* A customized version of how to manage
 * an entropy pool. The code has been tested and works.
 * 
 * Written September 7, 2025
 * Updated September 8, 2025
 * Updated September 9, 2025
 * Updated July 20, 2026
 * by Dirk Mittler
 * 
 */

#include <linux/module.h>		/* Needed by all modules */
#include <linux/ktime.h>
#include <linux/spinlock.h>
#include <linux/crypto.h>
#include <crypto/hash.h>
#include <linux/scatterlist.h>
#include <linux/uaccess.h>
#include <linux/fs.h>
#include <linux/cdev.h>
#include <linux/device.h>
#include <linux/delay.h>
#include <linux/jiffies.h>
#include <linux/dcache.h>
#include <linux/path.h>
#include <linux/slab.h>
#include <linux/string.h>

#include <linux/init.h>
#include <linux/rtc.h>
#include <linux/dmi.h>

#include <linux/sched.h>
#include <linux/sched/signal.h>
#include <linux/irqflags.h>

#define MAX_IRQs 256
#define HMICROS_BEFORE_BYTES 255


/* The following are declarations of objects visible
 * to other modules.
 */

void add_interrupt_randomness(int irq);

int device_open(struct inode *node, struct file *device_file);
ssize_t device_read(struct file *device_file,
		char __user *dest_buffer, size_t count, loff_t *offset);
ssize_t device_write(struct file *device_file,
		const char __user *source_buffer, size_t count,
		loff_t *offset);
int device_release(struct inode *inode, struct file *file);
int __init load(void);
void __exit unload(void);

/* End of shared objects.
 */


static unsigned char entropy_out_buffer[32];
static u32 entropy_pool[18];

static s64 lit[MAX_IRQs];
static s64 ltbi[MAX_IRQs];

static u32 entropy_bit_addr;
static s32 fake_bytes;

static int entropy_user_pointer;

static DEFINE_SPINLOCK(entropy_spinlock);
static DEFINE_MUTEX(entropy_mutex);

static struct class *my_class;
static struct device *my_device;
static struct device *my_device_b;
static dev_t dev_num;
static dev_t dev_num_b;
static struct cdev my_cdev;
static struct cdev my_cdev_b;

static struct rtc_device *rtc;
static struct rtc_time tm;

pid_t controlling;


// This function reads the real-time clock...
static int rtc_reader(void)
{
    int err;

    // 1. Open the default hardware clock (usually "rtc0")
    rtc = rtc_class_open("rtc0");
    if (!rtc) {
        printk(KERN_ALERT "Failed to open rtc0.\n");
        return 0;
    }

    // 2. Read the time directly from the hardware
    err = rtc_read_time(rtc, &tm);
    if (err < 0) {
        pr_err("rtc_reader: Failed to read from hardware clock, error %d\n", err);
        rtc_class_close(rtc);
        return err;
    }

    // 3. Release the RTC device
    rtc_class_close(rtc);

    return 0;
}

// This function XORs the motherboard UUID into primary entropy pool.
static void uuid_reader(void)
{
    const char *uuid;
    char *out = (char *) entropy_pool;
	int i = 0;

    // Fetch the UUID string directly from the kernel's DMI cache.
    uuid = dmi_get_system_info(DMI_PRODUCT_UUID);

    // CRITICAL: Always check for NULL. Not all motherboards populate this!
    if (uuid) {
		while (i < 32 && *uuid) {
			while (*uuid < 48 || *uuid > 102 ||
					(*uuid > 57 && *uuid < 97)) {
				if (*uuid == 0)
					break;
				uuid++;
			}
			if (*uuid == 0)
				break;
			out[i++] ^= *uuid;
			uuid++;
		}
		
    } else {
        printk(KERN_ALERT "DMI_PRODUCT_UUID is not available on this hardware.\n");
    }

    return;
}

// Function calls rtc_reader() and hashes result into primary entropy pool.
static int rtc_hash(void) {
	int ret = 0;
	int failure = 0;
	struct crypto_shash *sha256_tfm;
	SHASH_DESC_ON_STACK(shash, sha256_tfm);
	struct scatterlist sg;

	// Generate data.
	ret = rtc_reader();
	if (ret) {
		printk(KERN_ERR "Failed to read 'rtc0'.\n");
		return ret;
	}

	// Data to be hashed
	const char *my_data = (char *) &tm;
	size_t data_len = sizeof(struct rtc_time);

	sha256_tfm = crypto_alloc_shash("sha256", 0, 0);
	if (IS_ERR(sha256_tfm)) {
		// Handle error
		printk(KERN_ERR "Failed to allocate SHA-256 transform\n");
		return -EFAULT;
	}
	
	// Set up the scatterlist
	sg_init_one(&sg, my_data, data_len);

	// Initialize the hash
	shash->tfm = sha256_tfm;
	ret = crypto_shash_init(shash);
	if (ret) {
		// Handle error
		printk(KERN_ERR "Failed to initialize SHA-256 hash\n");
		failure = -EFAULT;
		goto free_tfm;
	}

	// Update the hash with your data
	ret = crypto_shash_update(shash, my_data, data_len);
	if (ret) {
		// Handle error
		printk(KERN_ERR "Failed to update SHA-256 hash\n");
		failure = -EFAULT;
		goto free_tfm;
	}
	
	ret = crypto_shash_final(shash, (char *) entropy_pool);
	if (ret) {
		// Handle error
		printk(KERN_ERR "Failed to finalize SHA-256 hash\n");
		failure = -EFAULT;
	}
		// Else, we now have the hash in entropy_pool
	
	free_tfm:
    crypto_free_shash(sha256_tfm);
    
    return failure;
}


// This function XORs-in 1 bit.
static void randomize_bit(u32 *pool, u32 *bit_addr,
			int value, spinlock_t *lock) {
	int finger = 1;
	int word_addr = 0;
	int pos_within = 0;
	int temp = 0;

	if (!spin_trylock(lock))
		return;

	temp = *bit_addr;
	
	word_addr = temp & 0x000000e0;
	pos_within = temp & 0x0000001f;
	word_addr >>= 5;
	
	// It's possible that the CPU wouldn't malfunction
	// if told to left-shift by zero bits, but just in case...
	if (pos_within > 0)
		finger <<= pos_within;

	if (value)
		pool[word_addr] ^= finger;
	
	temp++;
	
	if (temp > 16383)			// 2**14 - 1.
		temp -= 4096;			// Power of 2.
	
	(*bit_addr) = temp;
	
	spin_unlock(lock);

}

// This function XORs-in 1 byte.
static void randomize_byte(u32 *pool, u32 *bit_addr,
			unsigned char value, spinlock_t *lock) {
	int finger = value;
	int word_addr = 0;
	int pos_within = 0;
	int temp = 0;
	
	if (!spin_trylock(lock))
		return;

	finger ^= (unsigned char) (jiffies << 5);
	temp = *bit_addr;
	
	word_addr = temp & 0x000000e0;
	pos_within = temp & 0x0000001f;
	word_addr >>= 5;
	
	if (pos_within > 24) {
		pos_within = 24;
		temp &= 0xfffffff8;
	}
	
	// It's possible that the CPU wouldn't malfunction
	// if told to left-shift by zero bits, but just in case...
	if (pos_within > 0)
		finger <<= pos_within;

	pool[word_addr] ^= finger;

	temp += 8 ;
	
	if (temp > 16383)			// 2**14 - 1.
		temp -= 4096;			// Power of 2.
	
	(*bit_addr) = temp;
	
	spin_unlock(lock);

}

void add_interrupt_randomness(int irq) {
    // ktime_t now;				// Used for 32-bit kernels.
    s64 cit = 0;
    s64 ctbi = 0;
    s64 diff = 0;
    
    if (irq >= MAX_IRQs || irq < 0)
		return;

    // Get the current time in ktime_t format
    // now = ktime_get();

    // Convert ktime_t to a 64-bit integer representing 0.5 uS units
    cit = ktime_get_ns() / 500;
    
    ctbi = cit - lit[irq];
    lit[irq] = cit;
	diff = ctbi - ltbi[irq];

#if CONFIG_HIGH_RES_TIMERS == 'y'
#else
diff /= (2000000 / CONFIG_HZ);
#endif

	if (diff > HMICROS_BEFORE_BYTES || -diff > HMICROS_BEFORE_BYTES) {
		ltbi[irq] = ctbi;
		randomize_byte(entropy_pool, &entropy_bit_addr,
						(unsigned char) diff, &entropy_spinlock);
		return;
	}
    if (diff > 1) {
		ltbi[irq] = ctbi;
		randomize_bit(entropy_pool, &entropy_bit_addr, 0,
						&entropy_spinlock);
		return;
	}
	if (diff < -1) {
		ltbi[irq] = ctbi;
		randomize_bit(entropy_pool, &entropy_bit_addr, 1,
						&entropy_spinlock);
		return;
	}
}

static int add_user_randomness(char byte) {
	u32 temp1 = 0;
	u32 temp2 = byte;
	int word_addr = entropy_user_pointer & 0x0000001c;
	int pos_within = entropy_user_pointer & 0x00000003;
	
	word_addr >>= 2;
	word_addr += 8;
	
	temp1 = entropy_pool[word_addr];
	
	if (pos_within > 0)
		temp2 <<= (pos_within << 3);
	
	temp1 ^= temp2;
	entropy_pool[word_addr] = temp1;
	
	entropy_user_pointer++;
	entropy_user_pointer &= 0x0000001f;
	
	fake_bytes++;
	if (fake_bytes > 65535)
		fake_bytes = 32;
	
	return 0;
}

// Function was bracketed by 'mutex_lock_interruptible()' elsewhere.
static int entropy_hash(void) {
	s64 *nonce = (s64 *) (entropy_pool + 16);
	int ret = 0;
	int failure = 0;
	struct crypto_shash *sha256_tfm;
	SHASH_DESC_ON_STACK(shash, sha256_tfm);
	struct scatterlist sg;

	// Data to be hashed.
	const char *my_data = (char *) (entropy_pool);
	size_t data_len = 72;

	(* nonce) ++;				// Increment nonce.
	
	sha256_tfm = crypto_alloc_shash("sha256", 0, 0);
	if (IS_ERR(sha256_tfm)) {
		// Handle error
		printk(KERN_ERR "Failed to allocate SHA-256 transform\n");
		return -EFAULT;
	}
	
	// Set up the scatterlist.
	sg_init_one(&sg, my_data, data_len);

	// Initialize the hash.
	shash->tfm = sha256_tfm;
	ret = crypto_shash_init(shash);
	if (ret) {
		// Handle error.
		printk(KERN_ERR "Failed to initialize SHA-256 hash\n");
		failure = -EFAULT;
		goto free_tfm;
	}

	// Update the hash with data.
	ret = crypto_shash_update(shash, my_data, data_len);
	if (ret) {
		// Handle error
		printk(KERN_ERR "Failed to update SHA-256 hash\n");
		failure = -EFAULT;
		goto free_tfm;
	}
	
	ret = crypto_shash_final(shash, (char *) entropy_out_buffer);
	if (ret) {
		// Handle error
		printk(KERN_ERR "Failed to finalize SHA-256 hash\n");
		failure = -EFAULT;
	}
		// Else, we now have the hash in entropy_out_buffer
	
	free_tfm:
    crypto_free_shash(sha256_tfm);
    
    return failure;
}


int device_open(struct inode *node, struct file *device_file) {
	return 0;
}

static int cmp_filename_from_file(struct file *file, char *string) {
    int ret = 0;
    char *buf, *path_buf;
    
    // Allocate buffer for the path
    buf = kmalloc(PATH_MAX, GFP_KERNEL);
    if (!buf)
        return 0;
    
    // Get the path string
    path_buf = d_path(&file->f_path, buf, PATH_MAX);
    
    // Check for errors
    if (IS_ERR(path_buf)) {
        kfree(buf);
        return 0;
    }
    
    // The path_buf now contains the full path including the filename
    ret = strcmp(path_buf, string);
    
    kfree(buf);
    
    return (ret == 0);
}

ssize_t device_read(struct file *device_file,
		char __user *dest_buffer, size_t count, loff_t *offset) {
	int i = 0;
	int ret1 = 0;
	long int ret2 = 0;
//	char randombyte = 0;

	int nonblock = device_file->f_flags & O_NONBLOCK;

	// The ability to read can be blocked until
	// the total entropy has reached 128 bits.
	if (cmp_filename_from_file(device_file, "/dev/specrandom_b")) {
		while((fake_bytes << 3) + entropy_bit_addr < 128) {
			if (nonblock)
				return -EAGAIN;
			if(msleep_interruptible(20))
				return -ERESTARTSYS;
		}
	}

	if (controlling == current->pid)
		goto skip_mutex_r;

	if (nonblock) {
		if (!mutex_trylock(&entropy_mutex)) {
			return -EAGAIN;
		}
	} else {
		if (mutex_lock_interruptible(&entropy_mutex)) {
			return -ERESTARTSYS;	// Sleep was interrupted by signal.
		}
	}
	
	controlling = current->pid;
	
	skip_mutex_r:

	if (count < 0) {
		controlling = 0;
		mutex_unlock(&entropy_mutex);
		printk(KERN_ERR "Negative num of random bytes requested.\n");
		return -EFAULT;
	}

	while (count > 0) {
		entropy_hash();
		
		i = count;
		if (i > 32)
			i = 32;
		
		count -= i;
		
		ret2 = copy_to_user(dest_buffer + ret1, entropy_pool, i);
		ret1 += i;
		
		if (ret2) {
			printk(KERN_ERR "Usermem fault.\n");
			count = -EFAULT;
			break;
		}
		
		*offset += i;
		
		if (signal_pending(current)) {
			controlling = 0;
			mutex_unlock(&entropy_mutex);
			return -EINTR;
		}

	}

	controlling = 0;
	mutex_unlock(&entropy_mutex);
	
	if (count < 0) {
		return count;
	} else {
		return ret1;
	}
}

ssize_t device_write(struct file *device_file,
		const char __user *source_buffer, size_t count,
		loff_t *offset) {
	int i = 0;
	char temp = 0;
	int ret1 = 0;
	int ret2 = 0;

	int nonblock = device_file->f_flags & O_NONBLOCK;

	if (controlling == current->pid)
		goto skip_mutex_w;

	if (nonblock) {
		if (!mutex_trylock(&entropy_mutex)) {
			return -EAGAIN;
		}
	} else {
		if (mutex_lock_interruptible(&entropy_mutex)) {
			return -ERESTARTSYS;	// Sleep was interrupted by signal.
		}
	}

	controlling = current->pid;
	
	skip_mutex_w:
	for ( ; i < count ; i++ ) {
		ret1 = copy_from_user(&temp, source_buffer + i, 1);
		if (ret1) {
			count = -EFAULT;
			break;
		}
		ret2 = add_user_randomness(temp);
		if (ret2) {
			count = ret2;
			break;
		}
		
		if (signal_pending(current)) {
			controlling = 0;
			mutex_unlock(&entropy_mutex);
			return -EINTR;
		}
	}

	if (count >= 0)
		*offset += count;

	controlling = 0;
	mutex_unlock(&entropy_mutex);
	
	return count;
}

int device_release(struct inode *inode, struct file *file) {
    // The kernel calls this when the last reference to the file is gone.
    // It's a good place to do cleanup, including force-unlocking if needed.
    if (mutex_is_locked(&entropy_mutex)) {
        mutex_unlock(&entropy_mutex);
    }
    controlling = 0;
    return 0;
}

const struct file_operations device_fops = {
	.owner = THIS_MODULE,
    .open = device_open,
    .read = device_read,
    .write = device_write,
    .release = device_release,
    // ... other handlers ...
};


int __init load(void) {
	int i = 0;
	int ret = 0;
	
	for ( ; i < MAX_IRQs ; i++ ) {
		lit[i] = 0;
		ltbi[i] = 0;
	}
	
	entropy_bit_addr = 0;
	fake_bytes = 0;
	entropy_user_pointer = 0;

	for ( i = 0 ; i < 18 ; i++ ) {
		entropy_pool[i] = 0;
	}
	
	for ( i = 0 ; i < 32 ; i++ ) {
		entropy_out_buffer[i] = 0;
	}
	
	ret = rtc_hash();				// Must come before uuid_reader().
	if (ret) {
		printk(KERN_ERR "Failed to initialize pool with rtc hash.\n");
		return ret;
	}
	uuid_reader();
	
    // Allocate two major numbers for device.
    ret = alloc_chrdev_region(&dev_num, 0, 1, "special_random");
    if (ret < 0) {
        printk(KERN_ALERT "Failed to allocate device number\n");
        return ret;
    }
    ret = alloc_chrdev_region(&dev_num_b, 0, 1, "special_random_b");
    if (ret < 0) {
        printk(KERN_ALERT "Failed to allocate device number _b\n");
        unregister_chrdev_region(dev_num, 1);
        return ret;
    }

    // Initialize the cdev structs and link to file_operations struct
    cdev_init(&my_cdev, &device_fops);
    my_cdev.owner = THIS_MODULE;
    cdev_init(&my_cdev_b, &device_fops);
    my_cdev_b.owner = THIS_MODULE;

    // Add the character devices to the system
    ret = cdev_add(&my_cdev, dev_num, 1);
    if (ret < 0) {
        printk(KERN_ALERT "Failed to add cdev\n");
        unregister_chrdev_region(dev_num, 1);
        unregister_chrdev_region(dev_num_b, 1);
        return ret;
    }
    ret = cdev_add(&my_cdev_b, dev_num_b, 1);
    if (ret < 0) {
        printk(KERN_ALERT "Failed to add cdev _b\n");
        cdev_del(&my_cdev);
        unregister_chrdev_region(dev_num, 1);
        unregister_chrdev_region(dev_num_b, 1);
        return ret;
    }
    
    // Create the device class and device files in /dev...
    
    my_class = class_create("dirks_class");
    if (IS_ERR(my_class)) {
        printk(KERN_ALERT "Failed to create device class\n");
        cdev_del(&my_cdev);
        cdev_del(&my_cdev_b);
        unregister_chrdev_region(dev_num, 1);
        unregister_chrdev_region(dev_num_b, 1);
        return PTR_ERR(my_class);
    }
    
    my_device = device_create(my_class, NULL, dev_num, NULL, "specrandom");
    if (IS_ERR(my_device)) {
        printk(KERN_ALERT "Failed to create device specrandom\n");
		class_destroy(my_class);
        cdev_del(&my_cdev);
        cdev_del(&my_cdev_b);
        unregister_chrdev_region(dev_num, 1);
        unregister_chrdev_region(dev_num_b, 1);
        return PTR_ERR(my_device);
    }

    my_device_b = device_create(my_class, NULL, dev_num_b, NULL, "specrandom_b");
    if (IS_ERR(my_device_b)) {
        printk(KERN_ALERT "Failed to create device specrandom_b\n");
		device_destroy(my_class, dev_num);
		class_destroy(my_class);
        cdev_del(&my_cdev);
        cdev_del(&my_cdev_b);
        unregister_chrdev_region(dev_num, 1);
        unregister_chrdev_region(dev_num_b, 1);
        return PTR_ERR(my_device_b);
    }

    printk(KERN_INFO "Module loaded, device files: /dev/specrandom(_b)\n");

	return 0;
}

void __exit unload(void) {
    device_destroy(my_class, dev_num);
    device_destroy(my_class, dev_num_b);
    class_destroy(my_class);
    cdev_del(&my_cdev);
    cdev_del(&my_cdev_b);
    unregister_chrdev_region(dev_num, 1);
    unregister_chrdev_region(dev_num_b, 1);
	printk(KERN_INFO "Module unloaded.\n");
}

module_init(load);
module_exit(unload);

// Already defined by kernel and critical
// to the computer working ...
// EXPORT_SYMBOL(add_interrupt_randomness);

MODULE_LICENSE("GPL");
MODULE_AUTHOR("Dirk Mittler");
MODULE_DESCRIPTION("A test project for entry pool management.");