#include <stdio.h>
#include <stdlib.h>

#include <ext2fs/ext2fs.h>

const char *usage = "[-l] [-c] [-L<n>] <fs> <dir>";
const char *prog;

struct options {
	int limit;
	int count;
	int level;
	const char *fsname;
	const char *dirname;
};

int modifies_fs(struct options *opt)
{
	return (opt->limit || opt->count);
}


void fuzz_dir(struct options *opt)
{
	int flags = 0;
	int offset;
	errcode_t err;
	ext2_filsys fs = 0;
	ext2_ino_t ino;
	struct ext2_inode inode;
	blk_t block, pblk;
	unsigned entry = 0;
	char *buffer;
	struct ext2_dx_entry *entries;
	struct ext2_dx_countlimit *limit;
	struct ext2_dx_root_info *root;

	if (modifies_fs(opt)) {
		flags |= EXT2_FLAG_RW;
	}

	err = ext2fs_open2(opt->fsname, NULL, flags, 0, 0, unix_io_manager, &fs);
	if (err) {
		com_err(prog, err, "(while opening '%s')", opt->fsname);
		return;
	}

	err = ext2fs_namei(fs, EXT2_ROOT_INO, EXT2_ROOT_INO, opt->dirname, &ino);
	if (err) {
        com_err(prog, err, "(while locating directory %s)", opt->dirname);
		goto out;
    }
	if (!ino) {
		com_err(prog, 0, "no directory %s", opt->dirname);
		goto out;
	}

	err = ext2fs_read_inode(fs, ino, &inode);
	if (err) {
		com_err(prog, 0, "(reading directory inode)");
		goto out;
	}

	if (!LINUX_S_ISDIR(inode.i_mode)) {
		com_err(prog, 0, "not a directory");
		goto out;
	}

	if ((inode.i_flags & EXT2_BTREE_FL) == 0) {
		com_err(prog, 0, "not a hash-indexed directory");
		goto out;
	}

	buffer = malloc(2 * fs->blocksize);
	if (!buffer) {
		com_err(prog, 0, "couldn't allocate htree buffer");
		goto out;
	}

	block = pblk = inode.i_block[0];
	err = io_channel_read_blk(fs->io, block, 1, buffer);
	if (err) {
		com_err(prog, err, "(reading dir block %u)", block);
		goto bufout;
	}

	root = (struct ext2_dx_root_info *) (buffer + 24);
	if (opt->level > root->indirect_levels) {
		com_err(prog, 0, "level too high");
		goto bufout;
	}

	entries = (struct ext2_dx_entry *) (buffer + 24 + root->info_length);
	limit = (struct ext2_dx_countlimit *) entries;

	while (opt->level--) {
		unsigned count = ext2fs_le16_to_cpu(limit->count);
		entry = (rand() % count);
		block = ext2fs_le32_to_cpu((entries + entry)->block) & 0x00ffffff;

		err = ext2fs_bmap(fs, ino, &inode, buffer, 0, block, &pblk);
		if (err) {
			com_err(prog, err, "(mapping logical block %u)", block);
			goto bufout;
		}

		err = io_channel_read_blk(fs->io, pblk, 1, buffer);
		if (err) {
			com_err(prog, err, "(reading dir block %u)", block);
			goto bufout;
		}

		entries = (struct ext2_dx_entry *) (buffer + 8);
		limit = (struct ext2_dx_countlimit *) entries;
	}

	if (opt->limit) {
		printf("Fuzzing limit of entry %u (block %u)\n", entry, block);
		printf("Old limit: %u\n", ext2fs_le16_to_cpu(limit->limit));
		limit->limit = rand();
		printf("New limit: %u\n", ext2fs_le16_to_cpu(limit->limit));
	} else {
		printf("Limit: %u\n", ext2fs_le16_to_cpu(limit->limit));
	}

	if (opt->count) {
		printf("Fuzzing count of entry %u (block %u)\n", entry, block);
		printf("Old count: %u\n", ext2fs_le16_to_cpu(limit->count));
		limit->count = rand();
		printf("New count: %u\n", ext2fs_le16_to_cpu(limit->count));
	} else {
		printf("Count: %u\n", ext2fs_le16_to_cpu(limit->count));
	}

	if (modifies_fs(opt)) {
		err = io_channel_write_blk(fs->io, pblk, 1, buffer);
		if (err) {
			com_err(prog, err, "(writing dir block)");
		}
	}

  bufout:
	free(buffer);

  out:
	err = ext2fs_close(fs);
	if (err) {
		com_err(prog, err, "(while closing '%s')", opt->fsname);
	}
}

int main(int argc, char *argv[])
{
	int err;
	struct options opt = {
		.limit = 0,
		.count = 0,
		.level = 0
	};
	char **arg;

	prog = argv[0];

	for (arg = argv + 1; *arg && (*arg)[0] == '-'; ++arg) {
		switch ((*arg)[1]) {
		case 'l':
			opt.limit = 1;
			break;
		case 'c':
			opt.count = 1;
			break;
		case 'L':
			if ((*arg)[2] >= '0' && (*arg)[2] <= '9') {
				opt.level = (*arg)[2] - '0';
				break;
			}
		default:
			com_err(prog, 0, "usage: %s", usage);
			exit(1);
		}
	}

	if (argc != (arg - argv) + 2) {
		com_err(prog, 0, "usage: %s", usage);
		exit(1);
	}

	srand(time(NULL));
	initialize_ext2_error_table();

	opt.fsname = arg[0];
	opt.dirname = arg[1];
	fuzz_dir(&opt);
	exit(err ? 1 : 0);
}