/*
 * 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, ...)
{
	char err[ERRMAX];
	va_list arg;

	va_start(arg, fmt);
	vsnprint(err, ERRMAX, fmt, arg);
	va_end(arg);
	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*(disk->s+1));
	mbr = malloc(secsize*disk->s);
	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 {
		memset(buf, 0, secsize*disk->s);
		if((sysfd = open(mbrfile, OREAD)) < 0)
			fatal("open %s: %r", mbrfile);
		if((nmbr = read(sysfd, buf, secsize*(disk->s+1))) < 0)
			fatal("read %s: %r", mbrfile);
		if(nmbr > secsize*disk->s)
			fatal("master boot record too large %d > %d", nmbr, secsize*disk->s);
		if(nmbr < secsize)
			nmbr = secsize;
		close(sysfd);
		memmove(buf+Toffset, mbr+Toffset, secsize-Toffset);
		memmove(mbr, buf, nmbr);
	}
	mbr[secsize-2] = 0x55;
	mbr[secsize-1] = 0xAA;
	nmbr = (nmbr+secsize-1)&~(secsize-1);
	if(seek(disk->wfd, 0, 0) < 0)
		fatal("seek to MBR sector: %r\n");
	if(write(disk->wfd, mbr, nmbr) != nmbr)
		fatal("writing MBR: %r");
	
	exits(0);
}
