#include "libisofs.h"

#include "eltorito.h"
#include "volume.h"
#include "util.h"
#include "messages.h"

#include <assert.h>
#include <malloc.h>
#include <time.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include <string.h>
#include <stdio.h>
#include <fcntl.h>


/**
 * This table should be written with accuracy values at offset
 * 8 of boot image, when used ISOLINUX boot loader
 */
struct boot_info_table {
	uint8_t bi_pvd				BP(1, 4);  /* LBA of primary volume descriptor */
	uint8_t bi_file				BP(5, 8);  /* LBA of boot file */
	uint8_t bi_length			BP(9, 12); /* Length of boot file */
	uint8_t bi_csum				BP(13, 16); /* Checksum of boot file */
	uint8_t bi_reserved			BP(17, 56); /* Reserved */	
};

struct partition_desc {
	uint8_t boot_ind;
	uint8_t begin_chs[3];
	uint8_t type;
	uint8_t end_chs[3];
	uint8_t start[4];
	uint8_t size[4];
};

struct hard_disc_mbr {
	uint8_t code_area[440];
	uint8_t opt_disk_sg[4];
	uint8_t pad[2];
	struct partition_desc partition[4];
	uint8_t sign1;
	uint8_t sign2;
};

static void
el_torito_image_free(struct el_torito_boot_image *image)
{
	assert(image);
	
	/* unref image node */
	iso_tree_free((struct iso_tree_node*)image->node);
	free(image);
}

static struct iso_tree_node_boot*
create_boot_image_node(struct iso_tree_node_file *image)
{
	struct iso_tree_node_boot *boot;
	struct iso_tree_node_dir *parent;
	
	assert(image);
	
	boot = calloc(1, sizeof(struct iso_tree_node_boot));
	boot->node.attrib = image->node.attrib;
	boot->node.type = LIBISO_NODE_BOOT;
	boot->node.name = strdup(image->node.name);
	boot->node.hide_flags = image->node.hide_flags;
	boot->node.refcount = 2; /* mine and tree */
	boot->loc.path = strdup(image->loc.path);
	boot->img = 1; /* it refers to image */
	
	/* replace iso_tree_node_file with the image */
	parent = image->node.parent;
	iso_tree_node_remove(parent, (struct iso_tree_node*)image);
	iso_tree_add_child(parent, (struct iso_tree_node*) boot);
	
	return boot;
}

static struct el_torito_boot_image *
create_image(const char *path, 
             enum eltorito_boot_media_type type)
{
	struct stat attrib;
	struct el_torito_boot_image *boot;
	int boot_media_type = 0;
	int load_sectors = 0; /* number of sector to load */
	unsigned char partition_type = 0;
	
	if (stat(path, &attrib)) {
		iso_msg_sorry(LIBISO_EL_TORITO_WRONG_IMG, "Can't access file.");
		libisofs_errno = NO_FILE;
		return NULL;
	}
	
	if ( !S_ISREG(attrib.st_mode) ) {
		iso_msg_sorry(LIBISO_EL_TORITO_WRONG_IMG, 
			"We can only create boot images from regular files");
		libisofs_errno = ELTORITO_WRONG_IMAGE;
		return NULL;
	}
	
	switch (type) {
	case ELTORITO_FLOPPY_EMUL:
		switch (attrib.st_size) {
		case 1200 * 1024:
			boot_media_type = 1; /* 1.2 meg diskette */
			break;
		case 1440 * 1024:
			boot_media_type = 2; /* 1.44 meg diskette */
			break;
		case 2880 * 1024:
			boot_media_type = 3; /* 2.88 meg diskette */
			break;
		default:
			{
			char msg[72];
			sprintf(msg, "Invalid image size %d Kb. Must be one of 1.2, 1.44"
			               "or 2.88 Mb", (int) attrib.st_size / 1024);
			iso_msg_sorry(LIBISO_EL_TORITO_WRONG_IMG, msg);
			libisofs_errno = ELTORITO_WRONG_IMAGE_SIZE;
			return NULL;
			}
			break;	
		}
		/* it seems that for floppy emulation we need to load 
		 * a single sector (512b) */
		load_sectors = 1;
		break;
	case ELTORITO_HARD_DISC_EMUL:
		{
		size_t i;
		int fd;
		struct hard_disc_mbr mbr;
		int used_partition;
		
		/* read the MBR on disc and get the type of the partition */
		fd = open(path, O_RDONLY);
		if ( fd == -1 ) {
			iso_msg_sorry(LIBISO_EL_TORITO_WRONG_IMG, "Can't open image file.");
			libisofs_errno = NO_READ_ACCESS;
			return NULL;
		}
		if ( read(fd, &mbr, sizeof(mbr)) ) {
			iso_msg_sorry(LIBISO_EL_TORITO_WRONG_IMG, 
			              "Can't read MBR from image file.");
			libisofs_errno = ELTORITO_WRONG_IMAGE;
			close(fd);
			return NULL;
		}
		close(fd);
		
		/* check valid MBR signature */
		if ( mbr.sign1 != 0x55 || mbr.sign2 != 0xAA ) {
			iso_msg_sorry(LIBISO_EL_TORITO_WRONG_IMG, 
			              "Invalid MBR. Wrong signature.");
			libisofs_errno = ELTORITO_WRONG_IMAGE;
			return NULL;
		}
		
		/* ensure single partition */
		used_partition = -1;
		for (i = 0; i < 4; ++i) {
			if (mbr.partition[i].type != 0) {
				/* it's an used partition */
				if (used_partition != -1) {
					char msg[72];
					sprintf(msg, "Invalid MBR. At least 2 partitions: %d and " 
									"%d, are being used\n", used_partition, i);
					iso_msg_sorry(LIBISO_EL_TORITO_WRONG_IMG, msg);
					libisofs_errno = ELTORITO_WRONG_IMAGE;
					return NULL;
				} else
					used_partition = i;
			}
		}
		partition_type = mbr.partition[used_partition].type;
		}
		boot_media_type = 4;
		
		/* only load the MBR */
		load_sectors = 1;
		break;
	case ELTORITO_NO_EMUL:
		boot_media_type = 0;
		break;	
	}
	
	boot = calloc(1, sizeof(struct el_torito_boot_image));
	boot->bootable = 1;
	boot->type = boot_media_type;
	boot->load_size = load_sectors;
	boot->partition_type = partition_type;
	return boot;
}

/* parent and name can be null */
static struct iso_tree_node_boot*
create_boot_catalog_node(struct iso_tree_node_dir *parent, 
				         const char *name)
{
	struct iso_tree_node_boot *boot;
	
	assert( (parent && name) || (!parent && !name) );
	
	boot = calloc(1, sizeof(struct iso_tree_node_boot));
	boot->node.attrib.st_mode = S_IFREG | 0444;
	boot->node.attrib.st_atime = 
		boot->node.attrib.st_mtime = 
		boot->node.attrib.st_ctime = time(NULL);
	boot->node.attrib.st_size = 2048;
	boot->node.type = LIBISO_NODE_BOOT;
	boot->node.name = name ? strdup(name) : NULL;
	boot->node.refcount = 1; /* mine */
	if (parent) {
		boot->node.refcount++; /* add tree ref */
		iso_tree_add_child(parent, (struct iso_tree_node*) boot);
	}
	return boot;
}

struct el_torito_boot_image*
iso_volume_set_boot_image(struct iso_volume *volume, 
                          struct iso_tree_node *image,
                          enum eltorito_boot_media_type type,
                          struct iso_tree_node_dir *dir, 
                          char *name)
{
	struct el_torito_boot_image *boot_image;
	struct iso_tree_node_boot *cat_node;
	struct el_torito_boot_catalog *catalog;
	struct iso_tree_node_file *imgfile;
	
	assert(volume && !volume->bootcat);
	assert(image && ISO_ISREG(image));
	assert(dir && name);
	
	imgfile = (struct iso_tree_node_file*)image;
	
	if (image->procedence == LIBISO_PREVIMG) {
		/* FIXME mmm, what to do here? */
		iso_msg_sorry(LIBISO_EL_TORITO_WRONG_IMG, "Can't use prev. image files "
			"as boot images");
		return NULL;
	}
	
	boot_image = create_image(imgfile->loc.path, type);
	if ( !boot_image ) 
		return NULL;
		
	/* create image node */
	boot_image->node = create_boot_image_node(imgfile);
	
	/* creates the catalog with the given default image */
	catalog = malloc(sizeof(struct el_torito_boot_catalog));
	if (!catalog) {
		el_torito_image_free(boot_image);
		return NULL;
	}
	catalog->image = boot_image;
	
	/* add catalog file */
	cat_node = create_boot_catalog_node(dir, name);
	catalog->node = cat_node;
	
	volume->bootcat = catalog;
	
	return boot_image;
}

struct el_torito_boot_image* 
iso_volume_set_boot_image_hidden(struct iso_volume *volume, const char* path,
                                 enum eltorito_boot_media_type type)
{
	struct el_torito_boot_image *boot_image;
	struct iso_tree_node_boot *cat_node;
	struct el_torito_boot_catalog *catalog;
	
	assert(volume && !volume->bootcat);
	assert(path);
	
	boot_image = create_image(path, type);
	if ( !boot_image ) 
		return NULL;
		
	/* create image node */
	{
	struct iso_tree_node_boot *boot;
	boot = calloc(1, sizeof(struct iso_tree_node_boot));
	/* we assume the stat won't fail, as we already stat the same file above */
	stat(path, &boot->node.attrib);
	boot->node.type = LIBISO_NODE_BOOT;
	boot->node.refcount = 1; /* mine */
	boot->loc.path = strdup(path);
	boot->img = 1; /* it refers to image */
	boot_image->node = boot;
	}
	
	/* creates the catalog with the given default image */
	catalog = malloc(sizeof(struct el_torito_boot_catalog));
	if (!catalog) {
		el_torito_image_free(boot_image);
		return NULL;
	}
	catalog->image = boot_image;
	
	/* add catalog file */
	cat_node = create_boot_catalog_node(NULL, NULL);
	catalog->node = cat_node;
	
	volume->bootcat = catalog;
	
	return boot_image;
}

struct el_torito_boot_image*
iso_volume_get_boot_image(struct iso_volume *volume,
                          struct iso_tree_node **imgnode,
                          struct iso_tree_node **catnode)
{
	struct el_torito_boot_catalog *catalog;
	
	assert(volume);
	
	catalog = volume->bootcat;
	if (!catalog)
		return NULL;
		
	if (imgnode)
		*imgnode = (struct iso_tree_node*)catalog->image->node;
	
	if (catnode)
		*catnode = (struct iso_tree_node*)catalog->node;
	
	return catalog->image;
}

void
iso_volume_remove_boot_image(struct iso_volume *volume)
{
	struct el_torito_boot_catalog *catalog;
	struct iso_tree_node *catnode;
	struct iso_tree_node *imgnode;
	assert(volume);
	
	catalog = volume->bootcat;
	if (!catalog)
		return;
	
	/* remove node from tree */
	catnode = (struct iso_tree_node*)catalog->node;
	if (catnode->parent)
		iso_tree_node_remove(catnode->parent, catnode);
	
	/* remove image from tree */
	imgnode = (struct iso_tree_node*)catalog->image->node;
	if (imgnode->parent)
		iso_tree_node_remove(imgnode->parent, imgnode);
	
	/* remove catalog */
	el_torito_boot_catalog_free(catalog);
	volume->bootcat = NULL;
}

void
el_torito_set_load_seg(struct el_torito_boot_image *bootimg, int segment)
{
	if (bootimg->type != ELTORITO_NO_EMUL)
		return;
	bootimg->load_seg = segment;
}

void
el_torito_set_load_size(struct el_torito_boot_image *bootimg, int sectors)
{
	if (bootimg->type != ELTORITO_NO_EMUL)
		return;
	bootimg->load_size = sectors;
}

void
el_torito_set_no_bootable(struct el_torito_boot_image *bootimg)
{
	bootimg->bootable = 0;
}

void
el_torito_patch_isolinux_image(struct el_torito_boot_image *bootimg)
{
	bootimg->isolinux = 1;
}

void el_torito_boot_catalog_free(struct el_torito_boot_catalog *cat)
{
	assert(cat);
	
	el_torito_image_free(cat->image);
	iso_tree_free((struct iso_tree_node*)cat->node);
	free(cat);
}

/**
 * Write the Boot Record Volume Descriptor
 */
void
el_torito_write_boot_vol_desc(struct ecma119_write_target *t, uint8_t *buf)
{
	struct el_torito_boot_catalog *cat = t->catalog;
	struct ecma119_boot_rec_vol_desc *vol = 
			(struct ecma119_boot_rec_vol_desc*)buf;
	
	assert(cat);
	
	vol->vol_desc_type[0] = 0;
	memcpy(vol->std_identifier, "CD001", 5);
	vol->vol_desc_version[0] = 1;
	memcpy(vol->boot_sys_id, "EL TORITO SPECIFICATION", 23);
	iso_lsb(vol->boot_catalog, t->catblock, 4);
}

static void 
write_validation_entry(struct ecma119_write_target *t, uint8_t *buf)
{
	size_t i;
	int checksum;
	
	struct el_torito_validation_entry *ve = 
		(struct el_torito_validation_entry*)buf;
	ve->header_id[0] = 1;
	ve->platform_id[0] = 0; /* 0: 80x86, 1: PowerPC, 2: Mac */
	ve->key_byte1[0] = 0x55;
	ve->key_byte2[0] = 0xAA;
	
	/* calculate the checksum, to ensure sum of all words is 0 */
	checksum = 0;
	for (i = 0; i < sizeof(struct el_torito_validation_entry); i += 2) {
		checksum -= buf[i];
		checksum -= (buf[i] << 8);
	}
	iso_lsb(ve->checksum, checksum, 2);
}

static void
patch_boot_image(uint8_t *buf, struct ecma119_write_target *t, ssize_t imgsize)
{
	struct boot_info_table *info;
	uint32_t checksum;
	ssize_t offset;
	
	memset(&info, 0, sizeof(info));
	
	/* compute checksum, as the the sum of all 32 bit words in boot image
	 * from offset 64 */
	checksum = 0;
	offset = (ssize_t) 64;
	
	while (offset < imgsize) {
		checksum += iso_read_lsb(buf + offset, 4);
		offset += 4;
	}
	if (offset != imgsize) {
		/* file length not multiple of 4 */
		iso_msg_warn(LIBISO_ISOLINUX_CANT_PATCH, 
			"Unexpected isolinux image length. Patch might not work.");
	}
	
	/* patch boot info table */
	info = (struct boot_info_table*)(buf + 8);
	memset(info, 0, sizeof(struct boot_info_table));
	iso_lsb(info->bi_pvd, t->ms_block + 16, 4);
	iso_lsb(info->bi_file, t->imgblock, 4);
	iso_lsb(info->bi_length, imgsize, 4);
	iso_lsb(info->bi_csum, checksum, 4);
}

static void
write_boot_image(uint8_t *buf, struct ecma119_write_target *t)
{
	ssize_t imgsize;
	struct el_torito_boot_image *image;
	
	image = t->catalog->image;
	imgsize= image->node->node.attrib.st_size;
	
	/* copy image contents to memory buffer */
	if (image->node->node.procedence == LIBISO_PREVIMG) {
		int block, nblocks;
		uint8_t *memloc;
		assert(t->src);
		
		block = 0;
		memloc = buf;
		nblocks = imgsize ? div_up(imgsize, 2048) : 1;
		while (block < nblocks) {
			if (t->src->read_block(t->src, 
			            image->node->loc.block + block, memloc)) {
				iso_msg_sorry(LIBISO_CANT_READ_IMG,
					"Unable to read boot image from previous image.");
				break; /* exit loop */
			}
			memloc += 2048;
			++block;
		}
	} else {
		int fd;
		ssize_t nread, tread;
		uint8_t *memloc;
		const char *path = image->node->loc.path;
		memloc = buf;
	
		fd = open(path, O_RDONLY);
		if (fd == -1) {
			iso_msg_sorry(LIBISO_CANT_READ_FILE, 
			       "Can't open boot image file for reading. Skipping");
			return;
		}
		
		tread = 0;
		while (tread < imgsize) {
			nread = read(fd, memloc, t->block_size);
			if (nread <= 0) {
				iso_msg_sorry(LIBISO_CANT_READ_FILE, 
			       "Problem reading boot image file. Skipping");
				break; /* exit loop */
			}
			tread += nread;
			memloc += nread;
		}
		close(fd);
	}
	
	if (image->isolinux) {
		/* we need to patch the image */
		patch_boot_image(buf, t, imgsize);
	}
}

/**
 * Write one section entry.
 * Currently this is used only for default image (the only supported just now)
 */
static void
write_section_entry(uint8_t *buf, struct ecma119_write_target *t)
{
	struct el_torito_boot_image *img;
	struct el_torito_section_entry *se = 
		(struct el_torito_section_entry*)buf;
		
	img = t->catalog->image;
	
	assert(img);

	se->boot_indicator[0] = img->bootable ? 0x88 : 0x00;
	se->boot_media_type[0] = img->type;
	iso_lsb(se->load_seg, img->load_seg, 2);
	se->system_type[0] = img->partition_type;
	iso_lsb(se->sec_count, img->load_size, 2);
	iso_lsb(se->block, t->imgblock, 4);
}

/**
 * Write El-Torito Boot Catalog plus image
 */
static void
write_catalog(struct ecma119_write_target *t, uint8_t *buf)
{
	struct el_torito_boot_catalog *cat = t->catalog;
	assert(cat);
	assert(cat->image);
	
	write_validation_entry(t, buf);
	
	/* write default entry */
	write_section_entry(buf + 32, t);

	write_boot_image(buf + 2048, t);
}

void
el_torito_wr_boot_vol_desc(struct ecma119_write_target *t, uint8_t *buf)
{
	assert(t->catalog);
	ecma119_start_chunking(t,
			       el_torito_write_boot_vol_desc,
			       2048,
			       buf);
}

void
el_torito_wr_catalog(struct ecma119_write_target *t, uint8_t *buf)
{
	off_t size;
	off_t imgsize;
	assert(t->catalog);
	
	/*
	 * Note that when this is called, we need to write el-torito info to
	 * image. Note, however, that the image could come from a previous session
	 * and maybe we don't know its size
	 */
	
	imgsize = t->catalog->image->node->node.attrib.st_size;
	if (imgsize == 0) {
		iso_msg_warn(LIBISO_EL_TORITO_BLIND_COPY,
			"Can get boot image size, only one block will be copied");
		imgsize = 2048;
	}
	
	size = 2048 + div_up(imgsize, t->block_size);
	
	ecma119_start_chunking(t,
			       write_catalog,
			       size,
			       buf);
}