//
// LUKS passphrase bruteforcer - ruby interface
//
// Licence: GPLv2
// Author: Y. Guillot (2009), based on work from cryptsetup1.1.0
//
// compile with:
//  gcc -O2 -o luksbf.so -lgcrypt -ldevmapper -shared -I path/to/ruby.h
// usage:
//  ruby -r luksbf -e 'p LuksBF.new("/dev/sda1").test("p4ssw0rd")'
// to distribute computation, distribute a script that uses #derive. Give
// it a file holding the 1st 512 bytes from the original encrypted device
// (ie the LUKS hdr), then on the original machine check the derived keys using
// #test_derived.
//

#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <stdint.h>
#include <unistd.h>
#include <fcntl.h>
#include <gcrypt.h>
#include <errno.h>
#include <sys/ioctl.h>
#include <linux/fs.h>
#include <netinet/in.h>	// ntohl
#include <libdevmapper.h>

#include <ruby.h>

#define fail(e) rb_raise(rb_eRuntimeError, e)

#define LUKS_CIPHERNAME_L 32
#define LUKS_CIPHERMODE_L 32
#define LUKS_HASHSPEC_L 32
#define LUKS_DIGESTSIZE 20 // since SHA1
#define LUKS_HMACSIZE 32
#define LUKS_SALTSIZE 32
#define LUKS_NUMKEYS 8
#define LUKS_MKD_ITER 10
#define LUKS_KEY_DISABLED 0x0000DEAD
#define LUKS_KEY_ENABLED  0x00AC71F3
#define LUKS_STRIPES 4000
#define LUKS_MAGIC "LUKS\xba\xbe"
#define LUKS_MAGIC_L 6
#define LUKS_PHDR_SIZE (sizeof(struct luks_phdr)/SECTOR_SIZE+1)
#define UUID_STRING_L 40
#define LUKS_ALIGN_KEYSLOTS 4096

struct luks_phdr {
	char		magic[LUKS_MAGIC_L];
	uint16_t	version;
	char		cipherName[LUKS_CIPHERNAME_L];
	char		cipherMode[LUKS_CIPHERMODE_L];
	char            hashSpec[LUKS_HASHSPEC_L];
	uint32_t	payloadOffset;
	uint32_t	keyBytes;
	char		mkDigest[LUKS_DIGESTSIZE];
	char		mkDigestSalt[LUKS_SALTSIZE];
	uint32_t	mkDigestIterations;
	char            uuid[UUID_STRING_L];

	struct {
		uint32_t active;

		/* parameters used for password processing */
		uint32_t passwordIterations;
		char     passwordSalt[LUKS_SALTSIZE];

		/* parameters used for AF store/load */
		uint32_t keyMaterialOffset;
		uint32_t stripes;
	} keyblock[LUKS_NUMKEYS];
};

struct luks_masterkey {
	size_t keyLength;
	char key[];
};

static void init_crypto(void)
{
	if (!gcry_control(GCRYCTL_INITIALIZATION_FINISHED_P)) {
		gcry_control(GCRYCTL_SUSPEND_SECMEM_WARN);
		gcry_control(GCRYCTL_INIT_SECMEM, 16384, 0);
		gcry_control(GCRYCTL_RESUME_SECMEM_WARN);
		gcry_control(GCRYCTL_INITIALIZATION_FINISHED, 0);
	}
}

struct luksbf_state {
	struct luks_phdr hdr;
	char *afkey[LUKS_NUMKEYS];
	int afkey_len[LUKS_NUMKEYS];
	char *params[LUKS_NUMKEYS];
	char *hexkey[LUKS_NUMKEYS];
	char *name;
	int device_sector_size;
	char tmpkey[128];
	char dmCipherSpec[128];
	char dev_uuid[128];
	char device[128];
	char fullpath[128];
};

static void setup_state(struct luksbf_state *state)
{
	struct luks_phdr *hdr = &state->hdr;

	if (memcmp(hdr->magic, LUKS_MAGIC, LUKS_MAGIC_L))
		fail("no LUKS magic");
	if (ntohs(hdr->version) != 1)
		fail("bad LUKS version");

	hdr->payloadOffset      = ntohl(hdr->payloadOffset);
	hdr->keyBytes           = ntohl(hdr->keyBytes);
	hdr->mkDigestIterations = ntohl(hdr->mkDigestIterations);

	sprintf(state->dev_uuid, "CRYPT-TEMP-temporary-cryptsetup-%d",getpid());
	state->name = state->dev_uuid + strlen("CRYPT-TEMP-");
	sprintf(state->fullpath, "%s/%s", dm_dir(), state->name);
	sprintf(state->dmCipherSpec, "%s-%s", hdr->cipherName, hdr->cipherMode);

	int i;
	for (i=0 ; i<LUKS_NUMKEYS ; ++i) {
		hdr->keyblock[i].active             = ntohl(hdr->keyblock[i].active);
		if (hdr->keyblock[i].active != LUKS_KEY_ENABLED)
			continue;

		hdr->keyblock[i].passwordIterations = ntohl(hdr->keyblock[i].passwordIterations);
		hdr->keyblock[i].keyMaterialOffset  = ntohl(hdr->keyblock[i].keyMaterialOffset);
		hdr->keyblock[i].stripes            = ntohl(hdr->keyblock[i].stripes);

		state->afkey_len[i] = hdr->keyblock[i].stripes * hdr->keyBytes;
		state->afkey[i] = malloc(state->afkey_len[i]);
		if (!state->afkey[i])
			fail("nomem for afkey");

		state->params[i] = malloc(hdr->keyBytes*2 + strlen(state->dmCipherSpec) + strlen(state->device) + 64);
		if (!state->params[i])
			fail("nomem for params");
		sprintf(state->params[i], "%s %*d 0 %s %u",
			state->dmCipherSpec, hdr->keyBytes*2, 0, state->device,
			hdr->keyblock[i].keyMaterialOffset);

		state->hexkey[i] = state->params[i] + 1 +
 			strlen(state->dmCipherSpec);
	}
}

#define MAX_PRF_BLOCK_LEN 80

static int pkcs5_pbkdf2(const char *hash,
		const char *P, size_t Plen, const char *S, size_t Slen,
		unsigned int c, char *DK, unsigned int dkLen)
{
	gcry_md_hd_t prf;
	static char U[MAX_PRF_BLOCK_LEN];
	static char T[MAX_PRF_BLOCK_LEN];
	static char tmp[LUKS_SALTSIZE+4];
	int PRF, i, k;
	unsigned int u, hLen, l, r;
	unsigned char *p;

	PRF = gcry_md_map_name(hash);
	if (!PRF)
		fail("gcry_md_map_name");

	hLen = gcry_md_get_algo_dlen(PRF);
	if (hLen == 0 || hLen > MAX_PRF_BLOCK_LEN)
		fail("gcry_md_get_algo_dlen");

	if (c == 0)
		fail("bad itercount");

	l = dkLen / hLen;
	if (dkLen % hLen)
		l++;
	r = dkLen - (l - 1) * hLen;

	if (gcry_md_open(&prf, PRF, GCRY_MD_FLAG_HMAC))
		fail("gcry_md_open");

	if (gcry_md_setkey(prf, P, Plen))
		fail("gcry_md_setkey");

	for (i = 1; (uint) i <= l; i++) {
		memset(T, 0, hLen);

		for (u = 1; u <= c ; u++) {
			gcry_md_reset(prf);

			if (u == 1) {
				memcpy(tmp, S, Slen);
				tmp[Slen + 0] = (i & 0xff000000) >> 24;
				tmp[Slen + 1] = (i & 0x00ff0000) >> 16;
				tmp[Slen + 2] = (i & 0x0000ff00) >> 8;
				tmp[Slen + 3] = (i & 0x000000ff) >> 0;

				gcry_md_write(prf, tmp, Slen+4);
			} else
				gcry_md_write(prf, U, hLen);

			p = gcry_md_read(prf, 0);
			if (!p)
				fail("gcry_md_read");
			memcpy(U, p, hLen);

			for (k=0 ; (uint)k<hLen ; k++)
				T[k] ^= U[k];
		}

		memcpy(DK + (i-1)*hLen, T, (uint)i == l ? r : hLen);
	}
	gcry_md_close(prf);
	return 0;
}

/*
// TODO don't use dmcrypt, do this offline
// but for this we would need to implement
// cbc-essiv:sha256
static int LUKS_decrypt_from_storage_nodm(char *dst, int dstlen,
		struct luks_phdr *hdr, char *dkey, int dkeylen)
{
	int mode, tg;
	gcry_cipher_hd_t cf;

	tg = gcry_cipher_map_name(hdr->cipherName);
	if (!tg)
		fail("unknown cipher");
	if (*(uint32_t*)hdr->cipherMode == 0x636263)	// 'cbc'	// TODO cbc-essid:sha256
		mode = GCRY_CIPHER_MODE_CBC;
	else
		fail("unknown cipher mode");

	gcry_cipher_open(&cf, tg, mode, 0);
	gcry_cipher_setkey(cf, dkey, dkeylen);
	//gcry_cipher_setiv(cf, dkey, dkeylen);
	//rb_funcall("@io", "read", dst, dstlen, offset);
	gcry_cipher_decrypt(cf, dst, dstlen, dst, dstlen);

	return 0;
}
*/

#include <libdevmapper.h>
#include <linux/dm-ioctl.h>

void dm_remove_device(const char *name)
{
	struct dm_task *dmt = NULL;
	dmt = dm_task_create(DM_DEVICE_REMOVE);
	if (!dmt)
		return;
	dm_task_set_name(dmt, name);
	dm_task_run(dmt);
	dm_task_destroy(dmt);
}

int dm_create_device(struct luksbf_state *state, int slot, uint64_t size)
{
	struct dm_task *dmt = NULL;
	struct dm_info dmi;
	int r = -EINVAL;

	//printf("%llu | %s | %s | %s\n", size, state->params[slot], state->name, state->dev_uuid);
	if (!(dmt = dm_task_create(DM_DEVICE_CREATE)))
		fail("task_create");
	if (!dm_task_set_name(dmt, state->name))
		fail("task_set_name");
	if (!dm_task_set_uuid(dmt, state->dev_uuid))
		fail("task_set_uuid");
	if (!dm_task_set_ro(dmt))
		fail("task_set_ro");
	if (!dm_task_add_target(dmt, 0, size, "crypt", state->params[slot]))
		fail(state->params[slot]);
	if (!dm_task_run(dmt))
		fail("task_run");
	if (!dm_task_get_info(dmt, &dmi)) {
		dm_remove_device(state->name);
		fail("task_get_info");
	}
	r = 0;

	dm_task_destroy(dmt);
	dm_task_update_nodes();
	return r;
}


#define SECTOR_SHIFT           9
#define SECTOR_SIZE            (1 << SECTOR_SHIFT)

int LUKS_decrypt_from_storage(struct luksbf_state *state, int slot)
{
	int size, devfd;

	size = (state->afkey_len[slot] + state->device_sector_size - 1) /
		state->device_sector_size * state->device_sector_size /
		SECTOR_SIZE;

	if (dm_create_device(state, slot, size) < 0)
		fail("setup_mapping (check syslog)");

	devfd = open(state->fullpath, O_RDONLY);
	if (devfd == -1) {
		dm_remove_device(state->name);
		fail("open temporary keystore device");
	}

	if (read(devfd, state->afkey[slot], state->afkey_len[slot]) < 0) {
		close(devfd);
		dm_remove_device(state->name);
		fail("failed to access temporary keystore device");
	}

	close(devfd);
	dm_remove_device(state->name);
	return 0;
}

static int LUKS_derive_key(struct luksbf_state *state, int slot,
		char *key, int keylen)
{
	int r;
	struct luks_phdr *hdr = &state->hdr;

	r = pkcs5_pbkdf2(hdr->hashSpec, key, keylen,
			hdr->keyblock[slot].passwordSalt, LUKS_SALTSIZE,
			hdr->keyblock[slot].passwordIterations, state->tmpkey,
			hdr->keyBytes);

	unsigned i;
	for (i=0 ; i<state->hdr.keyBytes ; i++)
		sprintf(&state->hexkey[slot][i * 2], "%02x",
				(unsigned char)state->tmpkey[i]);
	state->hexkey[slot][i*2] = ' ';
	return r;
}

static void XORblock(char *dst, char *src, size_t n)
{
	unsigned i;
	for (i=0 ; i<n ; i++)
		dst[i] ^= src[i];
}

static void hash_buf(char *buf, uint32_t iv, int len, gcry_md_hd_t h)
{
	iv = htonl(iv);
	gcry_md_reset(h);
	gcry_md_write(h, (unsigned char*)&iv, sizeof(iv));
	gcry_md_write(h, buf, len);
	memcpy(buf, gcry_md_read(h, 0), len);
}

static void diffuse(char *buf, size_t sz, int dsz, gcry_md_hd_t h)
{
	int b, p, i;
	b = sz/dsz;
	p = sz%dsz;

	for (i=0 ; i<b ; i++)
		hash_buf(buf+dsz*i, i, dsz, h);
	if (p)
		hash_buf(buf+dsz*i, i, p,   h);
}

static int AF_merge(struct luksbf_state *state, int slot)
{
	int HID, dsz;
	unsigned i;
	int sz = state->hdr.keyBytes;
	gcry_md_hd_t md;


	HID = gcry_md_map_name(state->hdr.hashSpec);
	if (!HID)
		fail("bad hash spec");
	dsz = gcry_md_get_algo_dlen(HID);
	if (gcry_md_open(&md, HID, 0))
		fail("md_open");

	memset(state->tmpkey, 0, sizeof(state->tmpkey));
	for (i=0 ; i<state->hdr.keyblock[slot].stripes-1 ; i++) {
		XORblock(state->tmpkey, state->afkey[slot]+sz*i, sz);
		diffuse(state->tmpkey, sz, dsz, md);
	}
	XORblock(state->tmpkey, state->afkey[slot]+sz*i, sz);
	return 0;
}

static int LUKS_verify_master_key(struct luksbf_state *state)
{
	char buf[LUKS_DIGESTSIZE];
	pkcs5_pbkdf2(state->hdr.hashSpec, state->tmpkey, state->hdr.keyBytes,
			state->hdr.mkDigestSalt, LUKS_SALTSIZE,
			state->hdr.mkDigestIterations, buf, LUKS_DIGESTSIZE);

	if (memcmp(buf, state->hdr.mkDigest, LUKS_DIGESTSIZE))
		return -1;
	return 0;
}

static int LUKS_open_key(struct luksbf_state *state, int slot)
{
	if (LUKS_decrypt_from_storage(state, slot))
		fail("decrypt_from_storage");

	if (AF_merge(state, slot))
		fail("AF_merge");

	if (LUKS_verify_master_key(state))
		return -1;

	return 0;
}

static int test_key(struct luksbf_state *state, char *key, int keylen)
{
	int i;

	for (i=0 ; i<LUKS_NUMKEYS ; i++) {
		if (LUKS_derive_key(state, i, key, keylen) < 0)
			continue;

		if (!LUKS_open_key(state, i))
			return i;
	}

	return -1;
}

#define RSTATE(o) ((struct luksbf_state *)DATA_PTR(o))
#define iv_get_state(o) RSTATE(rb_ivar_get(o, rb_intern("@state")))

/* test a pass on all keyslots, returns the one that is unlocked */
static VALUE rb_test(VALUE self, VALUE arg)
{
	int ret;
	arg = rb_obj_as_string(arg);
	ret = test_key(iv_get_state(self), RSTRING_PTR(arg), RSTRING_LEN(arg));
	if (ret < 0)
		return Qnil;
	return rb_int2inum(ret);
}

/* test a precomputed derived key for a specific slot */
static VALUE rb_test_derived(VALUE self, VALUE dkey_o, VALUE slot_o)
{
	int slot = rb_num2ulong(slot_o);
	struct luksbf_state *state = iv_get_state(self);

	dkey_o = rb_obj_as_string(dkey_o);
	if (RSTRING_LEN(dkey_o) != (int)state->hdr.keyBytes*2)
		fail("invalid derived hex key");
	memcpy(state->hexkey[slot], RSTRING_PTR(dkey_o), RSTRING_LEN(dkey_o));
	return LUKS_open_key(state, slot) ? Qnil : Qtrue;
}

/* derive a key from a passphrase and a keyslot */
static VALUE rb_derive(VALUE self, VALUE pass_o, VALUE slot_o)
{
	struct luksbf_state *state = iv_get_state(self);
	int slot = rb_num2ulong(slot_o);

	pass_o = rb_obj_as_string(pass_o);
	if (LUKS_derive_key(state, slot, RSTRING_PTR(pass_o), RSTRING_LEN(pass_o)))
		return Qnil;
	return rb_str_new(state->hexkey[slot], state->hdr.keyBytes*2);
}

/* bf_state destructor */
static void free_state(void *ptr)
{
	struct luksbf_state *state = (struct luksbf_state *)ptr;
	int i;
	for (i=0 ; i<LUKS_NUMKEYS ; i++) {
		if (state->afkey[i])
			free(state->afkey[i]);
		if (state->params[i])
			free(state->params[i]);
	}
	free(state);
}

/* initialize the object from a path to a disk image/device */
static VALUE rb_init(VALUE self, VALUE arg)
{
	int fd;
	struct luksbf_state *state = ruby_xmalloc(sizeof(struct luksbf_state));
	memset(state, 0, sizeof(*state));
	VALUE state_o = rb_data_object_alloc(0, state, 0, free_state);

	arg = rb_obj_as_string(arg);
	if (rb_funcall(rb_cFile, rb_intern("exists?"), 1, arg) != Qtrue)
		rb_raise(rb_eArgError, "Invalid path");

	memcpy(state->device, RSTRING_PTR(arg), RSTRING_LEN(arg));

	fd = open(state->device, O_RDONLY);
	if (fd == -1 || read(fd, &state->hdr, sizeof(struct luks_phdr)) <
			(int)sizeof(struct luks_phdr))
		fail("cannot read device");
	ioctl(fd, BLKSSZGET, &state->device_sector_size);
	close(fd);

	setup_state(state);

	rb_ivar_set(self, rb_intern("@state"), state_o);

	return self;
}

static VALUE luksbf;

void Init_luksbf(void)
{
	init_crypto();

	luksbf = rb_define_class("LuksBF", rb_cObject);
	rb_define_method(luksbf, "initialize", rb_init, 1);
	rb_define_method(luksbf, "test", rb_test, 1);
	rb_define_method(luksbf, "test_derived", rb_test_derived, 2);
	rb_define_method(luksbf, "derive", rb_derive, 2);
}