/*	$NetBSD: installboot.c,v 1.8 2009/03/18 10:22:34 cegger Exp $	*/

/*
 * Copyright (c) 2000 NONAKA Kimihiro (nonaka@NetBSD.org).
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 * 3. The name of the author may not be used to endorse or promote products
 *    derived from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
 * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
 * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
 * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include <sys/param.h>
#include <sys/exec_elf.h>
#include <sys/bootblock.h>

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <fcntl.h>
#include <err.h>
#include <string.h>

int nowrite, verbose;
char *boot, *dev;

void usage(void);
int devread(int, void *, daddr_t, size_t, char *);
char *load_boot(char *, size_t *);
int load_prep_partition(int, struct mbr_partition *);
int main(int, char **);

void
usage(void)
{

	fprintf(stderr, "usage: %s [-n] [-v] <boot> <device>\n",
	    getprogname());
	exit(1);
}

int
devread(int fd, void *buf, daddr_t blk, size_t size, char *msg)
{

	if (lseek(fd, (off_t)dbtob(blk), SEEK_SET) != dbtob(blk)) {
		warn("%s: devread: lseek", msg);
		return 1;
	}
	if (read(fd, buf, size) != size) {
		warn("%s: devread: read", msg);
		return 1;
	}
	return 0;
}

char *
load_boot(char *boot, size_t *bootsize)
{
	Elf32_Ehdr eh;
	Elf32_Phdr ph;
	struct stat st;
	int fd;
	int i;
	size_t imgsz = 0;
	char *bp = NULL;

	if ((fd = open(boot, O_RDONLY)) < 0) {
		warn("open: %s", boot);
		return NULL;
	}

	if (fstat(fd, &st) != 0) {
		warn("fstat: %s", boot);
		goto out;
	}

	/*
	 * First, check ELF.
	 */
	if (read(fd, &eh, sizeof(eh)) != sizeof(eh)) {
		warn("read: eh: %s", boot);
		goto out;
	}
	if (memcmp(eh.e_ident, ELFMAG, SELFMAG) != 0 ||
	    eh.e_ident[EI_CLASS] != ELFCLASS32) {
		lseek(fd, 0L, SEEK_SET);
		goto notelf;
	}
	if (be16toh(eh.e_machine) != EM_PPC) {
		warn("not PowerPC binary.");
		goto out;
	}

	for (i = 0; i < be16toh(eh.e_phnum); i++) {
		(void)lseek(fd, be32toh(eh.e_phoff) + sizeof(ph) * i, SEEK_SET);
		if (read(fd, &ph, sizeof(ph)) != sizeof(ph)) {
			warn("read: ph: %s", boot);
			goto out;
		}

		if ((be32toh(ph.p_type) != PT_LOAD) ||
		    !(be32toh(ph.p_flags) & PF_X))
			continue;

		imgsz = st.st_size - be32toh(ph.p_offset);
		lseek(fd, be32toh(ph.p_offset), SEEK_SET);
		break;
	}

notelf:
	/*
	 * Second, check PReP bootable image.
	 */
	if (imgsz == 0) {
		char buf[DEV_BSIZE];

		printf("Bootable image: ");
		if (load_prep_partition(fd, 0)) {
			warn("no PReP bootable image.");
			goto out;
		}

		if (lseek(fd, (off_t)dbtob(1), SEEK_SET) != dbtob(1)) {
			warn("bootable image lseek sector 1");
			goto out;
		}
		if (read(fd, buf, DEV_BSIZE) != DEV_BSIZE) {
			warn("read: start/size");
			goto out;
		}

		imgsz = le32toh(*(u_int32_t *)(buf + sizeof(u_int32_t)))
		    - dbtob(2);
		lseek(fd, le32toh(*(u_int32_t *)buf), SEEK_SET);
	}

	if ((bp = (char *)calloc(roundup(imgsz, DEV_BSIZE), 1)) == NULL) {
		warn("calloc: no memory for boot image.");
		goto out;
	}

	if (read(fd, bp, imgsz) != imgsz) {
		warn("read: boot image: %s", boot);
		goto out;
	}

	if (verbose) {
		printf("image size = %d\n", imgsz);
	}

	*bootsize = roundup(imgsz, DEV_BSIZE);

	close(fd);
	return bp;

out:
	if (bp != NULL)
		free(bp);
	if (fd >= 0)
		close(fd);
	return NULL;
}

int
load_prep_partition(int devfd, struct mbr_partition *ppp)
{
	char mbr[512];
	struct mbr_partition *mbrp;
	int i;

	if (devread(devfd, mbr, MBR_BBSECTOR, DEV_BSIZE, "MBR") != 0)
		return 1;
	if (*(u_int16_t *)&mbr[MBR_MAGIC_OFFSET] != htole16(MBR_MAGIC)) {
		warn("no MBR_MAGIC");
		return 1;
	}

	mbrp = (struct mbr_partition *)&mbr[MBR_PART_OFFSET];
	for (i = 0; i < MBR_PART_COUNT; i++) {
		if (mbrp[i].mbrp_type == MBR_PTYPE_PREP)
			break;
	}
	if (i == MBR_PART_COUNT) {
		warn("no PReP partition.");
		return 1;
	}

	if (verbose) {
		printf("PReP partition: start = %d, size = %d\n",
		    le32toh(mbrp[i].mbrp_start), le32toh(mbrp[i].mbrp_size));
	}

	if (ppp) {
		*ppp = mbrp[i];
		ppp->mbrp_start = le32toh(ppp->mbrp_start);
		ppp->mbrp_size = le32toh(ppp->mbrp_size);
	}

	return 0;
}

int
main(int argc, char **argv)
{
	struct mbr_partition ppp;
	size_t bootsize;
	int c;
	int boot00[512/sizeof(int)];
	int devfd = -1;
	char *bp;

	while ((c = getopt(argc, argv, "vn")) != EOF) {
		switch (c) {
		case 'n':
			nowrite = 1;
			break;
		case 'v':
			verbose = 1;
			break;
		default:
			usage();
			break;
		}
	}

	if (argc - optind < 2)
		usage();

	boot = argv[optind];
	dev = argv[optind + 1];
	if (verbose) {
		printf("boot: %s\n", boot);
		printf("dev: %s\n", dev);
	}

	if ((bp = load_boot(boot, &bootsize)) == NULL)
		return 1;

	if ((devfd = open(dev, O_RDONLY, 0)) < 0) {
		warn("open: %s", dev);
		goto out;
	}

	if (load_prep_partition(devfd, &ppp)) {
		warn("load_prep_partition");
		goto out;
	}

	if (bootsize + dbtob(2) > dbtob(ppp.mbrp_size)) {
		warn("boot image is too big.");
		goto out;
	}

	close(devfd);

	if (nowrite) {
		free(bp);
		return 0;
	}

	if ((devfd = open(dev, O_RDWR, 0)) < 0) {
		warn("open: %s", dev);
		goto out;
	}

	/*
	 * Write boot image.
	 */
	memset(boot00, 0, sizeof(boot00));
	(void)lseek(devfd, (off_t)dbtob(ppp.mbrp_start), SEEK_SET);
	if (write(devfd, boot00, sizeof(boot00)) != sizeof(boot00)) {
		warn("write boot00(prep mbr)");
		goto out;
	}

	(void)lseek(devfd, (off_t)dbtob(ppp.mbrp_start+1), SEEK_SET);
	boot00[0] = htole32(dbtob(2));
	boot00[1] = htole32(bootsize);
	if (write(devfd, boot00, sizeof(boot00)) != sizeof(boot00)) {
		warn("write boot00(prep start/size)");
		goto out;
	}

	if (devread(devfd, boot00, 1, DEV_BSIZE, "start/size") != 0)
		goto out;
	boot00[0] = htole32(dbtob(ppp.mbrp_start));
	boot00[1] = htole32(bootsize + dbtob(2));
	(void)lseek(devfd, (off_t)dbtob(1), SEEK_SET);
	if (write(devfd, boot00, sizeof(boot00)) != sizeof(boot00)) {
		warn("write boot00(master start/size)");
		goto out;
	}

	(void)lseek(devfd, (off_t)dbtob(ppp.mbrp_start+2), SEEK_SET);
	if (write(devfd, bp, bootsize) != bootsize) {
		warn("write boot loader");
		goto out;
	}

	close(devfd);
	free(bp);
	return 0;

out:
	if (devfd >= 0)
		close(devfd);
	if (bp != NULL)
		free(bp);
	return 1;
}