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

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

typedef struct {
	uchar	active;			/* active flag */
	uchar	starth;			/* starting head */
	uchar	starts;			/* starting sector */
	uchar	startc;			/* starting cylinder */
	uchar	type;			/* partition type */
	uchar	endh;			/* ending head */
	uchar	ends;			/* ending sector */
	uchar	endc;			/* ending cylinder */
	uchar	lba[4];			/* starting LBA */
	uchar	size[4];		/* size in sectors */
} Tentry;

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

	Type9	= 0x39,
};

/*
 * 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);
}

static void
putle32(void* v, u32int i)
{
	uchar *p;

	p = v;
	p[0] = i;
	p[1] = i>>8;
	p[2] = i>>16;
	p[3] = i>>24;
}

static void
writechs(Disk *disk, uchar *p, vlong lba)
{
	int c, h, s;

	s = lba % disk->s;
	h = (lba / disk->s) % disk->h;
	c = lba / (disk->s * disk->h);

	if(c >= 1024) {
		c = 1023;
		h = disk->h - 1;
		s = disk->s - 1;
	}

	p[0] = h;
	p[1] = ((s+1) & 0x3F) | ((c>>2) & 0xC0);
	p[2] = c;
}

static void
wrtentry(Disk *disk, Tentry *tp, int type, u32int base, u32int lba, u32int end)
{
	tp->active = 0x80;		/* make this sole partition active */
	tp->type = type;
	writechs(disk, &tp->starth, lba);
	writechs(disk, &tp->endh, end-1);
	putle32(tp->lba, lba-base);
	putle32(tp->size, end-lba);
}

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

	flag9 = 0;
	mbrfile = nil;
	ARGBEGIN {
	case '9':
		flag9 = 1;
		break;
	case 'm':
		mbrfile = EARGF(usage());
		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");
	/*
	 * we need to cope with 4k-byte sectors on some newer disks.
	 * we're only interested in 512 bytes of mbr, so
	 * on 4k disks, rely on /dev/sd to read-modify-write.
	 */
	secsize = 512;
	if(disk->secsize != secsize)
		fprint(2, "%s: sector size %lld not %ld, should be okay\n",
			argv0, disk->secsize, 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);
	}

	if(flag9){
		tp = (Tentry*)(mbr+Toffset);
		memset(tp, 0, secsize-Toffset);
		wrtentry(disk, tp, Type9, 0, disk->s, disk->secs);
	}
	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);
}