/* $NetBSD: mkubootimage.c,v 1.18 2014/09/30 10:21:50 msaitoh Exp $ */

/*-
 * Copyright (c) 2010 Jared D. McNeill <jmcneill@invisible.ca>
 * 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. 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 <u.h>
#include <libc.h>
#include <flate.h>
#include "uboot.h"

#define __arraycount nelem

#define CRCPOLY 0xEDB88320
static ulong *crctab;

/* ripped from ape libc */
u32int
ntohl(u32int x)
{
	u32int n;
	unsigned char *p;

	n = x;
	p = (unsigned char*)&n;
	return (p[0]<<24)|(p[1]<<16)|(p[2]<<8)|p[3];
}

u32int
htonl(u32int h)
{
	u32int n;
	unsigned char *p;

	p = (unsigned char*)&n;
	p[0] = h>>24;
	p[1] = h>>16;
	p[2] = h>>8;
	p[3] = h;
	return n;
}

/* ripped from NetBSD */
u32int
bswap32(u32int x)
{
	return	((x << 24) & 0xff000000 ) |
		((x <<  8) & 0x00ff0000 ) |
		((x >>  8) & 0x0000ff00 ) |
		((x >> 24) & 0x000000ff );
}

static enum uboot_image_os image_os = IH_OS_PLAN9;
static enum uboot_image_arch image_arch = IH_ARCH_UNKNOWN;
static enum uboot_image_type image_type = IH_TYPE_UNKNOWN;
static enum uboot_image_comp image_comp = IH_COMP_NONE;
static u32int image_loadaddr = 0;
static u32int image_entrypoint = 0;
static char *image_name;
static u32int image_magic = IH_MAGIC;

static const struct uboot_os {
	enum uboot_image_os os;
	const char *name;
} uboot_os[] = {
	{ IH_OS_OPENBSD,	"openbsd" },
	{ IH_OS_NETBSD,		"netbsd" },
	{ IH_OS_FREEBSD,	"freebsd" },
	{ IH_OS_LINUX,		"linux" },
	{ IH_OS_PLAN9,		"plan9" },
};

static enum uboot_image_os
get_os(const char *name)
{
	unsigned int i;

	for (i = 0; i < __arraycount(uboot_os); i++) {
		if (strcmp(uboot_os[i].name, name) == 0)
			return uboot_os[i].os;
	}

	return IH_OS_UNKNOWN;
}

static const char *
get_os_name(enum uboot_image_os os)
{
	unsigned int i;

	for (i = 0; i < __arraycount(uboot_os); i++) {
		if (uboot_os[i].os == os)
			return uboot_os[i].name;
	}

	return "Unknown";
}

static const struct uboot_arch {
	enum uboot_image_arch arch;
	const char *name;
} uboot_arch[] = {
	{ IH_ARCH_ARM,		"arm" },
	{ IH_ARCH_I386,		"i386" },
	{ IH_ARCH_MIPS,		"mips" },
	{ IH_ARCH_MIPS64,	"mips64" },
	{ IH_ARCH_PPC,		"powerpc" },
	{ IH_ARCH_OPENRISC,	"or1k" },
	{ IH_ARCH_ARM64,	"arm64" },
};

static enum uboot_image_arch
get_arch(const char *name)
{
	unsigned int i;

	for (i = 0; i < __arraycount(uboot_arch); i++) {
		if (strcmp(uboot_arch[i].name, name) == 0)
			return uboot_arch[i].arch;
	}

	return IH_ARCH_UNKNOWN;
}

static const char * 
get_arch_name(enum uboot_image_arch arch)
{
	unsigned int i;

	for (i = 0; i < __arraycount(uboot_arch); i++) {
		if (uboot_arch[i].arch == arch)
			return uboot_arch[i].name;
	}

	return "Unknown";
}

static const struct uboot_type {
	enum uboot_image_type type;
	const char *name;
} uboot_type[] = {
	{ IH_TYPE_STANDALONE,	"standalone" },
	{ IH_TYPE_KERNEL,	"kernel" },
	{ IH_TYPE_RAMDISK,	"ramdisk" },
	{ IH_TYPE_FILESYSTEM,	"fs" },
	{ IH_TYPE_SCRIPT,	"script" },
};

static enum uboot_image_type
get_type(const char *name)
{
	unsigned int i;

	for (i = 0; i < __arraycount(uboot_type); i++) {
		if (strcmp(uboot_type[i].name, name) == 0)
			return uboot_type[i].type;
	}

	return IH_TYPE_UNKNOWN;
}

static const char *
get_type_name(enum uboot_image_type type)
{
	unsigned int i;

	for (i = 0; i < __arraycount(uboot_type); i++) {
		if (uboot_type[i].type == type)
			return uboot_type[i].name;
	}

	return "Unknown";
}

static const struct uboot_comp {
	enum uboot_image_comp comp;
	const char *name;
} uboot_comp[] = {
	{ IH_COMP_NONE,		"none" },
	{ IH_COMP_GZIP,		"gz" },
	{ IH_COMP_BZIP2,	"bz2" },
	{ IH_COMP_LZMA,		"lzma" },
	{ IH_COMP_LZO,		"lzo" },
};

static enum uboot_image_comp
get_comp(const char *name)
{
	unsigned int i;

	for (i = 0; i < __arraycount(uboot_comp); i++) {
		if (strcmp(uboot_comp[i].name, name) == 0)
			return uboot_comp[i].comp;
	}

	return IH_TYPE_UNKNOWN;
}

static const char *
get_comp_name(enum uboot_image_comp comp)
{
	unsigned int i;

	for (i = 0; i < __arraycount(uboot_comp); i++) {
		if (uboot_comp[i].comp == comp)
			return uboot_comp[i].name;
	}

	return "Unknown";
}

static void
usage(void)
{
	fprint(2, "usage: mkubootimage -A "
	    "<arm|arm64|i386|mips|mips64|or1k|powerpc>");
	fprint(2, " -C <none|bz2|gz|lzma|lzo>");
	fprint(2, " -O <openbsd|netbsd|freebsd|linux|plan9>");
	fprint(2, " -T <standalone|kernel|ramdisk|fs|script>");
	fprint(2, " -a <addr> [-e <ep>] [-m <magic>] -n <name>");
	fprint(2, " <srcfile> <dstfile>\n");

	exits("usage");
}

static void
dump_header(struct uboot_image_header *hdr)
{
	long tm = ntohl(hdr->ih_time);

	print(" magic:       0x%08ux\n", ntohl(hdr->ih_magic));
	print(" time:        %s", ctime(tm));
	print(" size:        %ud\n", ntohl(hdr->ih_size));
	print(" load addr:   0x%08ux\n", ntohl(hdr->ih_load));
	print(" entry point: 0x%08ux\n", ntohl(hdr->ih_ep));
	print(" data crc:    0x%08ux\n", ntohl(hdr->ih_dcrc));
	print(" os:          %d (%s)\n", hdr->ih_os, get_os_name(hdr->ih_os));
	print(" arch:        %d (%s)\n", hdr->ih_arch, get_arch_name(hdr->ih_arch));
	print(" type:        %d (%s)\n", hdr->ih_type, get_type_name(hdr->ih_type));
	print(" comp:        %d (%s)\n", hdr->ih_comp, get_comp_name(hdr->ih_comp));
	print(" name:        %s\n", (char*)hdr->ih_name);
	print(" header crc:  0x%08ux\n", ntohl(hdr->ih_hcrc));
}

static int
generate_header(struct uboot_image_header *hdr, int kernel_fd)
{
	u8int *p;
	Dir *st;
	ulong crc;
	u32int dsize, size_buf[2];
	long n;

	st = dirfstat(kernel_fd);
	if(st == nil){
		perror("stat");
		return 1;
	}

	if(st->length < 8192){
		sysfatal("kernel too small");
	}

	if(st->length + sizeof(*hdr) > (u32int)-1){
		fprint(2, "fatal: kernel too big\n");
		return 1;
	}

	p = mallocz(st->length, 1);
	if(p == nil){
		perror("malloc");
		return 1;
	}

	/* XXX: do blockwise reads */
	n = read(kernel_fd, p, st->length);
	if(n != st->length)
		sysfatal("short read: %r");

	if (image_type == IH_TYPE_SCRIPT) {
		dsize = st->length + (sizeof(u32int) * 2);
		size_buf[0] = htonl(st->length);
		size_buf[1] = htonl(0);
		crc = blockcrc(crctab, 0, size_buf, sizeof(u32int) * 2);
		crc = blockcrc(crctab, crc, p, st->length);
	} else {
		dsize = st->length;
		crc = blockcrc(crctab, 0, p, st->length);
	}

	memset(hdr, 0, sizeof(*hdr));
	hdr->ih_magic = htonl(image_magic);
	hdr->ih_time = htonl(st->mtime);
	hdr->ih_size = htonl(dsize);
	hdr->ih_load = htonl(image_loadaddr);
	hdr->ih_ep = htonl(image_entrypoint);
	hdr->ih_dcrc = htonl(crc);
	hdr->ih_os = image_os;
	hdr->ih_arch = image_arch;
	hdr->ih_type = image_type;
	hdr->ih_comp = image_comp;
	strncpy((char *)hdr->ih_name, image_name, sizeof(hdr->ih_name)-1);
	crc = blockcrc(crctab, 0, (void *)hdr, sizeof(*hdr));
	hdr->ih_hcrc = htonl(crc);

	dump_header(hdr);

	return 0;
}

static int
write_image(struct uboot_image_header *hdr, int kernel_fd, int image_fd)
{
	u8int buf[4096];
	long rlen, wlen;
	Dir *st;
	u32int size_buf[2];

	st = dirfstat(kernel_fd);
	if(st == nil){
		perror("stat");
		return 1;
	}

	wlen = write(image_fd, hdr, sizeof(*hdr));
	if (wlen != sizeof(*hdr))
		sysfatal("short write: %r");

	if (image_type == IH_TYPE_SCRIPT) {
		size_buf[0] = htonl(st->length);
		size_buf[1] = htonl(0);
		wlen = write(image_fd, &size_buf[0], sizeof(size_buf));
		if (wlen != sizeof(size_buf))
			sysfatal("short write: %r");
	}

	seek(kernel_fd, 0, 0);
	while ((rlen = read(kernel_fd, buf, sizeof(buf))) > 0) {
		wlen = write(image_fd, buf, rlen);
		if (wlen != rlen)
			sysfatal("short write: %r");
	}

	if(rlen < 0)
		sysfatal("read kernel: %r");

	return 0;
}

void
main(int argc, char *argv[])
{
	struct uboot_image_header hdr;
	const char *src, *dest, *arg;
	char *ep;
	int kernel_fd, image_fd;
	unsigned long long num;

	crctab = mkcrctab(CRCPOLY);

	ARGBEGIN{
	case 'A':	/* arch */
		image_arch = get_arch(EARGF(usage()));
		break;
	case 'C':	/* comp */
		image_comp = get_comp(EARGF(usage()));
		break;
	case 'O':	/* os */
		image_os = get_os(EARGF(usage()));
		break;
	case 'T':	/* type */
		image_type = get_type(EARGF(usage()));
		break;
	case 'a':	/* addr */
		arg = EARGF(usage());
		num = strtoull((arg), &ep, 0);
		if (*ep != '\0' || num == 0 ||
		    ((signed long long)num != (long)num &&
		     num != (u32int)num))
			sysfatal("illegal number -- %s", arg);
		image_loadaddr = (u32int)num;
		break;
	case 'E':	/* ep (byte swapped) */
	case 'e':	/* ep */
		arg = EARGF(usage());
		num = strtoull(arg, &ep, 0);
		if (*ep != '\0' ||  num == 0 ||
		    ((signed long long)num != (long)num &&
		     num != (u32int)num))
			sysfatal("illegal number -- %s", arg);
		image_entrypoint = (u32int)num;
		if (ARGC() == 'E')
			image_entrypoint = bswap32(image_entrypoint);
		break;
	case 'm':	/* magic */
		arg = EARGF(usage());
		num = strtoul(arg, &ep, 0);
		if (*ep != '\0' || num == 0)
			sysfatal("illegal number -- %s", arg);
		image_magic = (u32int)num;
	case 'n':	/* name */
		image_name = strdup(EARGF(usage()));
		break;
	case 'h':
	default:
		usage();
		/* NOTREACHED */
	}ARGEND

	if (argc != 2)
		usage();

	if (image_entrypoint == 0)
		image_entrypoint = image_loadaddr;

	if (image_arch == IH_ARCH_UNKNOWN ||
	    image_type == IH_TYPE_UNKNOWN ||
	    (image_type != IH_TYPE_SCRIPT && image_loadaddr == 0) ||
	    image_name == nil)
		usage();

	src = argv[0];
	dest = argv[1];

	kernel_fd = open(src, OREAD);
	if (kernel_fd == -1) {
		perror("open kernel");
		exits("open");
	}
	image_fd = create(dest, OWRITE|OTRUNC, 0666);
	if (image_fd == -1) {
		perror("open image");
		exits("open");
	}

	if (generate_header(&hdr, kernel_fd) != 0)
		exits("generate_header");

	if (write_image(&hdr, kernel_fd, image_fd) != 0)
		exits("write_image");

	close(image_fd);
	close(kernel_fd);

	exits(nil);
}