/*
 * install new master boot record boot code on PC disk.
 */

#include <u.h>
#include <libc.h>
#include <disk.h>

enum {
	Toffset = 0x1BE	/* offset of partition table */
};

/*
 * Default boot block prints an error message and reboots. 
 */
static int ndefmbr = Toffset;
static char defmbr[512] = {
[0x000]	0xEB, 0x3C, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
	0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
[0x03E] 0xFA, 0xFC, 0x8C, 0xC8, 0x8E, 0xD8, 0x8E, 0xD0,
	0xBC, 0x00, 0x7C, 0xBE, 0x77, 0x7C, 0xE8, 0x19,
	0x00, 0x33, 0xC0, 0xCD, 0x16, 0xBB, 0x40, 0x00,
	0x8E, 0xC3, 0xBB, 0x72, 0x00, 0xB8, 0x34, 0x12,
	0x26, 0x89, 0x07, 0xEA, 0x00, 0x00, 0xFF, 0xFF,
	0xEB, 0xD6, 0xAC, 0x0A, 0xC0, 0x74, 0x09, 0xB4,
	0x0E, 0xBB, 0x07, 0x00, 0xCD, 0x10, 0xEB, 0xF2,
	0xC3,  'N',  'o',  't',  ' ',  'a',  ' ',  'b',
	 'o',  'o',  't',  'a',  'b',  'l',  'e',  ' ',
	 'd',  'i',  's',  'c',  ' ',  'o',  'r',  ' ',
	 'd',  'i',  's',  'c',  ' ',  'e',  'r',  'r',
	 'o',  'r', '\r', '\n',  'P',  'r',  'e',  's',
	 's',  ' ',  'a',  'l',  'm',  'o',  's',  't',
	 ' ',  'a',  'n',  'y',  ' ',  'k',  'e',  'y',
	 ' ',  't',  'o',  ' ',  'r',  'e',  'b',  'o',
	 'o',  't',  '.',  '.',  '.', 0x00, 0x00, 0x00,
};

void
usage(void)
{
	fprint(2, "usage: disk/mbr [-m mbrfile] disk\n");
	exits("usage");
}

void
fatal(char *fmt, ...)
{
	int n;
	char err[128];
	va_list arg;

	va_start(arg, fmt);
	n = doprint(err, err+sizeof(err), fmt, arg) - err;
	va_end(arg);
	err[n] = 0;
	fprint(2, "mbr: %s\n", err);
	exits(err);
}

void
main(int argc, char **argv)
{
	Disk *disk;
	uchar *mbr, *buf;
	char *mbrfile;
	ulong secsize;
	int sysfd, nmbr;

	mbrfile = nil;
	ARGBEGIN {
	case 'm':
		mbrfile = ARGF();
		break;
	default:
		usage();
	} ARGEND

	if(argc < 1)
		usage();

	disk = opendisk(argv[0], 0, 0);
	if(disk == nil)
		fatal("opendisk %s: %r", argv[0]);

	if(disk->type == Tfloppy)
		fatal("will not install mbr on floppy");
	if(disk->secsize != 512)
		fatal("secsize %d invalid", disk->secsize);

	secsize = disk->secsize;
	buf = malloc(secsize);
	mbr = malloc(secsize);
	if(buf == nil || mbr == nil)
		fatal("out of memory");

	/*
	 * Start with initial sector from disk.
	 */
	if(seek(disk->fd, 0, 0) < 0)
		fatal("seek to boot sector: %r\n");
	if(read(disk->fd, mbr, secsize) != secsize)
		fatal("reading boot sector: %r");

	if(mbrfile == nil){
		nmbr = ndefmbr;
		memmove(mbr, defmbr, nmbr);
	} else {
		if((sysfd = open(mbrfile, OREAD)) < 0)
			fatal("open %s: %r", mbrfile);
		if((nmbr = read(sysfd, buf, secsize)) < 0)
			fatal("read %s: %r", mbrfile);
		if(nmbr > Toffset)
			fatal("master boot record too large");
		close(sysfd);
		memmove(mbr, buf, nmbr);
	}

	if(nmbr < Toffset)
		memset(mbr+nmbr, 0, Toffset-nmbr);

	mbr[secsize-2] = 0x55;
	mbr[secsize-1] = 0xAA;

	if(seek(disk->wfd, 0, 0) < 0)
		fatal("seek to MBR sector: %r\n");
	if(write(disk->wfd, mbr, secsize) != secsize)
		fatal("writing MBR: %r");
	
	exits(0);
}