/*
 * SqueamishOssifrage.c
 * CSAW 2011 CTF Challenge Kernel Module
 * Jon Oberheide <jon@oberheide.org>
 *
 * Decrypt the flag...and win?
 *
 * Information:
 *
 *   /proc/csaw/key:     get/set symmetric key
 *   /proc/csaw/encrypt: get/set encrypted payload
 *   /proc/csaw/decrypt: get/set decrypted payload
 *
 * Example usage:
 *
 *   # set a key
 *   $ echo "this is a key" > /proc/csaw/key
 *
 *   # encrypt payload with that key
 *   $ echo "hello world" > /proc/csaw/encrypt
 *
 *   # retrieve encrypted payload
 *   $ cat /proc/csaw/encrypt > payload.enc
 *
 *   # decrypt encrypted payload
 *   $ cat payload.enc > /proc/csaw/decrypt
 *
 *   # retrieve decrypted payload
 *   $ cat /proc/csaw/decrypt > payload.dec
 *
 *   # view decrypted payload
 *   $ cat payload.dec
 *   hello world
 */

#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/init.h>
#include <linux/proc_fs.h>
#include <linux/string.h>
#include <asm/uaccess.h>

#define PROC_CSAW    "csaw"
#define PROC_KEY     "key"
#define PROC_ENCRYPT "encrypt"
#define PROC_DECRYPT "decrypt"
#define MAX_LENGTH   128
#define MAX_ROUNDS   64
#define DELTA        0x9e3779b9

MODULE_LICENSE("GPL");
MODULE_AUTHOR("Jon Oberheide");
MODULE_DESCRIPTION("CSAW 2011 CTF Challenge Kernel Module");

static struct proc_dir_entry *proc_csaw;
static struct proc_dir_entry *proc_key;
static struct proc_dir_entry *proc_decrypt;
static struct proc_dir_entry *proc_encrypt;

static uint32_t key[4] = { 0, 0, 0, 0 };
static uint32_t rounds = 32;
static unsigned char enc_blob[MAX_LENGTH];
static unsigned char dec_blob[MAX_LENGTH];

void encrypt(uint32_t k[], char *buf, int len, int rounds) __attribute__ ((optimize (0)));
void decrypt(uint32_t k[], char *buf, int len, int rounds) __attribute__ ((optimize (0)));
void scramble(uint32_t k[], uint32_t t[], uint32_t *s, uint32_t *d) __attribute__ ((optimize (0)));
void descramble(uint32_t k[], uint32_t t[], uint32_t *s, uint32_t *d) __attribute__ ((optimize (0)));

void
scramble(uint32_t k[], uint32_t t[], uint32_t *s, uint32_t *d)
{
	uint32_t n, sbox[8];
	uint32_t y = t[0], z = t[1];
	uint32_t delta = DELTA, sum = 0;

	for (n = 0; n < 8; n++) {
		sbox[n] = key[n % 4];
	}
	for (n = 0; n < 32; n++) {
		sum += delta;
		y += ((z << 4) + k[0]) ^ (z+sum) ^ ((z >> 5) + k[1]);
		z += ((y << 4) + k[2]) ^ (y+sum) ^ ((y >> 5) + k[3]);
	}
	t[0] = y; s = d + n; t[1] = z;
}

void
encrypt(uint32_t k[], char *buf, int len, int rounds)
{
	int i;

	if (rounds >= 0) {
		if (rounds % 3 == 0) {
			k[0] ^= k[1]; k[1] ^= k[0]; k[0] ^= k[1];
			k[2] ^= k[3]; k[3] ^= k[2]; k[2] ^= k[3];
		} else if (rounds % 3 == 1) {
			k[0] ^= k[2]; k[2] ^= k[0]; k[0] ^= k[2];
			k[1] ^= k[3]; k[3] ^= k[1]; k[1] ^= k[3];
		} else if (rounds % 3 == 2) {
			k[0] ^= k[3]; k[3] ^= k[0]; k[0] ^= k[3];
			k[2] ^= k[1]; k[1] ^= k[2]; k[2] ^= k[1];
		}
		encrypt(k, buf, len, rounds-1);
	} else {
		for (i = 0; i < len / 8; i++) {
			scramble(k, (uint32_t *)(buf + (i * 8)), (uint32_t *) 0, (uint32_t *) DELTA);
		}
	}
}

void
descramble(uint32_t k[], uint32_t t[], uint32_t *s, uint32_t *d)
{
	uint32_t n, sbox[8];
	uint32_t y = t[0], z = t[1];
	uint32_t delta = DELTA, sum = delta << 5;

	for (n = 0; n < 8; n++) {
		sbox[n] = key[n % 4];
	}
	for (n = 0; n < 32; n++) {
		z -= ((y << 4) + k[2]) ^ (y + sum) ^ ((y >> 5) + k[3]);
		y -= ((z << 4) + k[0]) ^ (z + sum) ^ ((z >> 5) + k[1]);
		sum -= delta;
	}
	t[0] = y; s = d + n; t[1] = z;
}

void
decrypt(uint32_t k[], char *buf, int len, int rounds)
{
	int i;

	if (rounds >= 0) {
		if (rounds % 3 == 0) {
			k[0] ^= k[1]; k[1] ^= k[0]; k[0] ^= k[1];
			k[2] ^= k[3]; k[3] ^= k[2]; k[2] ^= k[3];
		} else if (rounds % 3 == 1) {
			k[0] ^= k[2]; k[2] ^= k[0]; k[0] ^= k[2];
			k[1] ^= k[3]; k[3] ^= k[1]; k[1] ^= k[3];
		} else if (rounds % 3 == 2) {
			k[0] ^= k[3]; k[3] ^= k[0]; k[0] ^= k[3];
			k[2] ^= k[1]; k[1] ^= k[2]; k[2] ^= k[1];
		}
		decrypt(k, buf, len, rounds-1);
	} else {
		for (i = 0; i < len / 8; i++) {
			descramble(k, (uint32_t *)(buf + (i * 8)), (uint32_t *) 0, (uint32_t *) DELTA);
		}
	}
}

int
key_write(struct file *file, const char __user *ubuf, unsigned long count, void *data)
{
	char buf[MAX_LENGTH];

	printk(KERN_INFO "csaw: called key_write\n");

	memset(buf, 0, sizeof(buf));

	if (count > MAX_LENGTH) {
		count = MAX_LENGTH;
	}
	if (copy_from_user(&buf, ubuf, count)) {
		return -EFAULT;
	}

	sscanf(buf, "%16c\t%d", (char *) key, &rounds);

	return count;
}

int
key_read(char *page, char **start, off_t off, int count, int *eof, void *data)
{
	printk(KERN_INFO "csaw: called key_read\n");

	memcpy(page, key, sizeof(key));

	return sizeof(key);
}

int
encrypt_write(struct file *file, const char __user *ubuf, unsigned long count, void *data)
{
	char buf[MAX_LENGTH];
	uint32_t session[4];

	printk(KERN_INFO "csaw: called encrypt_write\n");

	memset(buf, 0, sizeof(buf));

	if (count > MAX_LENGTH) {
		count = MAX_LENGTH;
	}
	if (copy_from_user(&buf, ubuf, count)) {
		return -EFAULT;
	}

	memcpy(session, key, sizeof(key));

	encrypt(session, buf, sizeof(buf), rounds);

	memcpy(enc_blob, buf, sizeof(buf));

	return count;
}

int
encrypt_read(char *page, char **start, off_t off, int count, int *eof, void *data)
{
	printk(KERN_INFO "csaw: called encrypt_read\n");

	memcpy(page, enc_blob, sizeof(enc_blob));

	return sizeof(enc_blob);
}

int
decrypt_write(struct file *file, const char __user *ubuf, unsigned long count, void *data)
{
	char buf[MAX_LENGTH];
	uint32_t session[4];

	printk(KERN_INFO "csaw: called decrypt_write\n");

	memset(buf, 0, sizeof(buf));

	if (count > MAX_LENGTH) {
		count = MAX_LENGTH;
	}
	if (copy_from_user(&buf, ubuf, count)) {
		return -EFAULT;
	}

	memcpy(session, key, sizeof(key));

	decrypt(session, buf, sizeof(buf), rounds);

	memcpy(dec_blob, buf, sizeof(buf));

	return count;
}

int
decrypt_read(char *page, char **start, off_t off, int count, int *eof, void *data)
{
	printk(KERN_INFO "csaw: called decrypt_read\n");

	memcpy(page, dec_blob, sizeof(dec_blob));

	return sizeof(dec_blob);
}

int __init
csaw_init(void)
{
	printk(KERN_INFO "csaw: loading module\n");

	memset(enc_blob, 0, sizeof(enc_blob));
	memset(dec_blob, 0, sizeof(dec_blob));

	proc_csaw = proc_mkdir(PROC_CSAW, NULL);

	proc_key = create_proc_entry(PROC_KEY, 0666, proc_csaw);
	proc_key->read_proc = key_read;
	proc_key->write_proc = key_write;

	proc_encrypt = create_proc_entry(PROC_ENCRYPT, 0666, proc_csaw);
	proc_encrypt->read_proc = encrypt_read;
	proc_encrypt->write_proc = encrypt_write;

	proc_decrypt = create_proc_entry(PROC_DECRYPT, 0666, proc_csaw);
	proc_decrypt->read_proc = decrypt_read;
	proc_decrypt->write_proc = decrypt_write;

	printk(KERN_INFO "csaw: created /proc/csaw entries\n");

	return 0;
}

void __exit
csaw_exit(void)
{
	printk(KERN_INFO "csaw: unloading module\n");

	remove_proc_entry(PROC_KEY, proc_csaw);
	remove_proc_entry(PROC_ENCRYPT, proc_csaw);
	remove_proc_entry(PROC_DECRYPT, proc_csaw);
	remove_proc_entry(PROC_CSAW, NULL);

	printk(KERN_INFO "csaw: removed /proc/csaw entries\n");
}

module_init(csaw_init);
module_exit(csaw_exit);