/*
 * Copyright 2020-2021 Haiku, Inc. All rights reserved.
 * All rights reserved. Distributed under the terms of the MIT License.
 *
 * Authors:
 * 		Alexander von Gluck IV <kallisti5@unixzen.com>
 */


#include <errno.h>
#include <getopt.h>
#include <fcntl.h>
#include <unistd.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>

#define EXIT_FAILURE 1


//#define DEBUG_MBRTOOL
#ifdef DEBUG_MBRTOOL
#   define TRACE(x...) printf("mbrtool: " x)
#else
#   define TRACE(x...) ;
#endif

#define INFO(x...) printf("mbrtool: " x)


// Disk sector size, 512 assumed!
static const size_t kSectorSize = 512;

static void
print_usage(bool error)
{
	printf("\n");
	printf("usage: mbrtool (options) <diskImage> <id> <type> <start> <len>\n");
	printf("       <diskImage>            Disk image to operate on\n");
	printf("       <id>                   Partition ID (0-3)\n");
	printf("       <type>                 Partition type id (hex)\n");
	printf("       <start>                Partition start offset (KiB)\n");
	printf("       <len>                  Partition length (KiB)\n\n");
	printf("  Options:\n");
	printf("       -a, --active           Partition boot flag\n");
	printf("\nWarning: This tool requires precision!\n");
	printf("         Inputs are only lightly validated!\n\n");
	exit(error ? EXIT_FAILURE : 0);
}


static void
checkError(bool failed, const char *message)
{
	if (!failed)
		return;

	if (errno != 0)
		INFO("Error: %s: %s\n", message, strerror(errno));
	else
		INFO("Error: %s\n", message);

	exit(1);
}


static ssize_t
mbrWipe(int handle)
{
	// Blow away MBR while avoiding bootloader
	uint8_t emptyMBR[66] = {};
	// End of MBR marker
	emptyMBR[64] = 0x55;
	emptyMBR[65] = 0xAA;
	return pwrite(handle, emptyMBR, 66, 0x1BE);
}


static bool
mbrValid(int handle)
{
	// TODO: this is a really basic check and ignores invalid
	// partition table entries
	uint8_t mbrBytes[66] = {};
	ssize_t read = pread(handle, mbrBytes, 66, 0x1BE);
	checkError(read < 0, "failed to read MBR for validation");
	return (mbrBytes[64] == 0x55 && mbrBytes[65] == 0xAA);
}


static void
createPartition(int handle, int index, bool active, uint8_t type,
	uint64_t offset, uint64_t size)
{
	uint8_t bootable = active ? 0x80 : 0x0;
	uint8_t partition[16] = {
		bootable,		// bootable
		0xff, 0xff, 0xff,	// CHS first block (default to LBA)
		type,			// partition type
		0xff, 0xff, 0xff,	// CHS last block (default to LBA)
		0x00, 0x00, 0x00, 0x00,	// imageOffset in blocks (written below)
		0x00, 0x00, 0x00, 0x00	// imageSize in blocks (written below)
	};

	// fill in LBA values
	uint32_t partitionOffset = (uint32_t)(offset / kSectorSize);
	((uint32_t *)partition)[2] = partitionOffset;
	((uint32_t *)partition)[3] = (uint32_t)(size / kSectorSize);

	TRACE("%s: #%d %c bytes: %u-%u, sectors: %u-%u\n", __func__, index,
		active ? 'b' : '-', offset, offset + size, partitionOffset,
		partitionOffset + uint32_t(size / kSectorSize));

	ssize_t written = pwrite(handle, partition, 16, 512 - 2 - 16 * (4 - index));
	checkError(written != 16, "failed to write partition entry");

	if (active) {
		// make it bootable
		written = pwrite(handle, &partitionOffset, 4, offset + 512 - 2 - 4);
		checkError(written != 4, "failed to make partition bootable");
	}
	return;
}


int
main(int argc, char *argv[])
{
	const char *imageFile = NULL;

	int partType = -1;
	int partIndex = -1;
	int64_t partStartOffset = -1;
	int64_t partLength = -1;
	bool partBootable = false;

	while (1) {
		int c;
		static struct option long_options[] = {
			{"active", no_argument, 0, 'a'},
			{"help", no_argument, 0, 'h'},
			{0, 0, 0, 0}
		};

		opterr = 1; /* don't print errors */
		c = getopt_long(argc, argv, "+ha", long_options, NULL);

		if (c == -1)
			break;

		switch (c) {
			case 's':
				print_usage(false);
				break;
			case 'a':
				partBootable = true;
				break;
			default:
				print_usage(true);
		}
	}

	if ((argc - optind) != 5)
		print_usage(true);

	for (int index = optind; index < argc; index++) {
		if (imageFile == NULL)
			imageFile = argv[index];
		else if (partIndex < 0)
			partIndex = atoi(argv[index]);
		else if (partType < 0)
			partType = (int)strtol(argv[index], NULL, 0);
		else if (partStartOffset < 0 ) {
			partStartOffset = (int64_t)strtol(argv[index], NULL, 10);
			partStartOffset *= 1024;
		} else if (partLength < 0) {
			partLength = (int64_t)strtol(argv[index], NULL, 10);
			partLength *= 1024;
		}
	}

	checkError(partIndex < 0 || partIndex > 3,
		"invalid partition index, valid range is 0-3");
	checkError(partType < 0,  "Incorrect Partition Type!");

	checkError((partStartOffset + partLength) > 2089072000000,
		"partitions beyond 2TiB are not accepted!");

	int imageFileHandle = open(imageFile, O_RDWR);
	checkError(imageFileHandle < 0, "failed to open disk image file");

	struct stat stat;
	int result = fstat(imageFileHandle, &stat);
	checkError(result != 0, "failed to stat image file");
	off_t imageSize = stat.st_size;

	if (!mbrValid(imageFileHandle)) {
		INFO("MBR of image is invalid, creating a fresh one.\n");
		mbrWipe(imageFileHandle);
	}

	// Just a warning. This is technically valid since MBR partition
	// definitions are purely within the first 512 bytes
	if (partStartOffset + partLength > imageSize)
		INFO("Warning: Partition extends beyond end of file!\n");

	createPartition(imageFileHandle, partIndex, partBootable, partType,
		partStartOffset, partLength);
	return 0;
}