/*
 * Copyright 2014 The Chromium OS Authors. All rights reserved.
 * Use of this source code is governed by a BSD-style license that can be
 * found in the LICENSE file.
 *
 * Verified boot kernel utility
 */

#include <errno.h>
#include <fcntl.h>
#include <getopt.h>
#include <inttypes.h>		/* For PRIu64 */
#include <linux/fs.h>		/* For BLKGETSIZE64 */
#include <stdarg.h>
#include <stdio.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/stat.h>
#include <unistd.h>

#include "futility.h"
#include "host_common.h"
#include "kernel_blob.h"
#include "vb1_helper.h"

static void Fatal(const char *format, ...)
{
	va_list ap;
	va_start(ap, format);
	fprintf(stderr, "ERROR: ");
	vfprintf(stderr, format, ap);
	va_end(ap);
	exit(1);
}

/* Global opts */
static int opt_verbose;
static int opt_vblockonly;
static uint64_t opt_pad = 65536;

/* Command line options */
enum {
	OPT_MODE_PACK = 1000,
	OPT_MODE_REPACK,
	OPT_MODE_VERIFY,
	OPT_ARCH,
	OPT_OLDBLOB,
	OPT_KLOADADDR,
	OPT_KEYBLOCK,
	OPT_SIGNPUBKEY,
	OPT_SIGNPRIVATE,
	OPT_VERSION,
	OPT_VMLINUZ,
	OPT_BOOTLOADER,
	OPT_CONFIG,
	OPT_VBLOCKONLY,
	OPT_PAD,
	OPT_VERBOSE,
	OPT_MINVERSION,
};

static const struct option long_opts[] = {
	{"pack", 1, 0, OPT_MODE_PACK},
	{"repack", 1, 0, OPT_MODE_REPACK},
	{"verify", 1, 0, OPT_MODE_VERIFY},
	{"arch", 1, 0, OPT_ARCH},
	{"oldblob", 1, 0, OPT_OLDBLOB},
	{"kloadaddr", 1, 0, OPT_KLOADADDR},
	{"keyblock", 1, 0, OPT_KEYBLOCK},
	{"signpubkey", 1, 0, OPT_SIGNPUBKEY},
	{"signprivate", 1, 0, OPT_SIGNPRIVATE},
	{"version", 1, 0, OPT_VERSION},
	{"minversion", 1, 0, OPT_MINVERSION},
	{"vmlinuz", 1, 0, OPT_VMLINUZ},
	{"bootloader", 1, 0, OPT_BOOTLOADER},
	{"config", 1, 0, OPT_CONFIG},
	{"vblockonly", 0, 0, OPT_VBLOCKONLY},
	{"pad", 1, 0, OPT_PAD},
	{"verbose", 0, &opt_verbose, 1},
	{"debug", 0, &debugging_enabled, 1},
	{NULL, 0, 0, 0}
};


static const char usage[] =
	"\n"
	"Usage:  " MYNAME " %s --pack <file> [PARAMETERS]\n"
	"\n"
	"  Required parameters:\n"
	"    --keyblock <file>         Key block in .keyblock format\n"
	"    --signprivate <file>      Private key to sign kernel data,\n"
	"                                in .vbprivk format\n"
	"    --version <number>        Kernel version\n"
	"    --vmlinuz <file>          Linux kernel bzImage file\n"
	"    --bootloader <file>       Bootloader stub\n"
	"    --config <file>           Command line file\n"
	"    --arch <arch>             Cpu architecture (default x86)\n"
	"\n"
	"  Optional:\n"
	"    --kloadaddr <address>     Assign kernel body load address\n"
	"    --pad <number>            Verification padding size in bytes\n"
	"    --vblockonly              Emit just the verification blob\n"
	"\nOR\n\n"
	"Usage:  " MYNAME " %s --repack <file> [PARAMETERS]\n"
	"\n"
	"  Required parameters:\n"
	"    --signprivate <file>      Private key to sign kernel data,\n"
	"                                in .vbprivk format\n"
	"    --oldblob <file>          Previously packed kernel blob\n"
	"                                (including verfication blob)\n"
	"\n"
	"  Optional:\n"
	"    --keyblock <file>         Key block in .keyblock format\n"
	"    --config <file>           New command line file\n"
	"    --version <number>        Kernel version\n"
	"    --kloadaddr <address>     Assign kernel body load address\n"
	"    --pad <number>            Verification blob size in bytes\n"
	"    --vblockonly              Emit just the verification blob\n"
	"\nOR\n\n"
	"Usage:  " MYNAME " %s --verify <file> [PARAMETERS]\n"
	"\n"
	"  Optional:\n"
	"    --signpubkey <file>"
	"       Public key to verify kernel keyblock,\n"
	"                                in .vbpubk format\n"
	"    --verbose                 Print a more detailed report\n"
	"    --keyblock <file>         Outputs the verified key block,\n"
	"                                in .keyblock format\n"
	"    --pad <number>            Verification padding size in bytes\n"
	"    --minversion <number>     Minimum combined kernel key version\n"
	"                              and kernel version\n"
	"\n";


/* Print help and return error */
static void print_help(const char *progname)
{
	printf(usage, progname, progname, progname);
}


/* Return an explanation when fread() fails. */
static const char *error_fread(FILE *fp)
{
	const char *retval = "beats me why";
	if (feof(fp))
		retval = "EOF";
	else if (ferror(fp))
		retval = strerror(errno);
	clearerr(fp);
	return retval;
}


/* This reads a complete kernel partition into a buffer */
static uint8_t *ReadOldKPartFromFileOrDie(const char *filename,
					 uint64_t *size_ptr)
{
	FILE *fp = NULL;
	struct stat statbuf;
	uint8_t *buf;
	uint64_t file_size = 0;

	if (0 != stat(filename, &statbuf))
		Fatal("Unable to stat %s: %s\n", filename, strerror(errno));

	if (S_ISBLK(statbuf.st_mode)) {
		int fd = open(filename, O_RDONLY);
		if (fd >= 0) {
			ioctl(fd, BLKGETSIZE64, &file_size);
			close(fd);
		}
	} else {
		file_size = statbuf.st_size;
	}
	Debug("%s size is 0x%" PRIx64 "\n", filename, file_size);
	if (file_size < opt_pad)
		Fatal("%s is too small to be a valid kernel blob\n");

	Debug("Reading %s\n", filename);
	fp = fopen(filename, "rb");
	if (!fp)
		Fatal("Unable to open file %s: %s\n", filename,
		      strerror(errno));

	buf = malloc(file_size);
	if (1 != fread(buf, file_size, 1, fp))
		Fatal("Unable to read entirety of %s: %s\n", filename,
		      error_fread(fp));

	if (size_ptr)
		*size_ptr = file_size;

	return buf;
}

/****************************************************************************/

static int do_vbutil_kernel2(int argc, char *argv[])
{
	char *filename = NULL;
	char *oldfile = NULL;
	char *keyblock_file = NULL;
	char *signpubkey_file = NULL;
	char *signprivkey_file = NULL;
	char *version_str = NULL;
	int version = -1;
	char *vmlinuz_file = NULL;
	char *bootloader_file = NULL;
	char *config_file = NULL;
	enum arch_t arch = ARCH_X86;
	uint64_t kernel_body_load_address = CROS_32BIT_ENTRY_ADDR;
	int mode = 0;
	int parse_error = 0;
	uint64_t min_version = 0;
	char *e;
	int i;
	int rv;
	VbKeyBlockHeader *keyblock = NULL;
	VbKeyBlockHeader *t_keyblock = NULL;
	VbPrivateKey *signpriv_key = NULL;
	VbPublicKey *signpub_key = NULL;
	uint8_t *kpart_data = NULL;
	uint64_t kpart_size = 0;
	uint8_t *vmlinuz_buf = NULL;
	uint64_t vmlinuz_size = 0;
	uint8_t *t_config_data;
	uint64_t t_config_size;
	uint8_t *t_bootloader_data;
	uint64_t t_bootloader_size;
	VbKernelPreambleHeader *preamble = NULL;
	uint8_t *kblob_data = NULL;
	uint64_t kblob_size = 0;
	uint8_t *vblock_data = NULL;
	uint64_t vblock_size = 0;

	while (((i = getopt_long(argc, argv, ":", long_opts, NULL)) != -1) &&
	       !parse_error) {
		switch (i) {
		default:
		case '?':
			/* Unhandled option */
			parse_error = 1;
			break;

		case 0:
			/* silently handled option */
			break;

		case OPT_MODE_PACK:
		case OPT_MODE_REPACK:
		case OPT_MODE_VERIFY:
			if (mode && (mode != i)) {
				fprintf(stderr,
					"Only one mode can be specified\n");
				parse_error = 1;
				break;
			}
			mode = i;
			filename = optarg;
			break;

		case OPT_ARCH:
			/* check the first 3 characters to also detect x86_64 */
			if ((!strncasecmp(optarg, "x86", 3)) ||
			    (!strcasecmp(optarg, "amd64")))
				arch = ARCH_X86;
			else if ((!strcasecmp(optarg, "arm")) ||
				 (!strcasecmp(optarg, "aarch64")))
				arch = ARCH_ARM;
			else if (!strcasecmp(optarg, "mips"))
				arch = ARCH_MIPS;
			else {
				fprintf(stderr,
					"Unknown architecture string: %s\n",
					optarg);
				parse_error = 1;
			}
			break;

		case OPT_OLDBLOB:
			oldfile = optarg;
			break;

		case OPT_KLOADADDR:
			kernel_body_load_address = strtoul(optarg, &e, 0);
			if (!*optarg || (e && *e)) {
				fprintf(stderr, "Invalid --kloadaddr\n");
				parse_error = 1;
			}
			break;

		case OPT_KEYBLOCK:
			keyblock_file = optarg;
			break;

		case OPT_SIGNPUBKEY:
			signpubkey_file = optarg;
			break;

		case OPT_SIGNPRIVATE:
			signprivkey_file = optarg;
			break;

		case OPT_VMLINUZ:
			vmlinuz_file = optarg;
			break;

		case OPT_BOOTLOADER:
			bootloader_file = optarg;
			break;

		case OPT_CONFIG:
			config_file = optarg;
			break;

		case OPT_VBLOCKONLY:
			opt_vblockonly = 1;
			break;

		case OPT_VERSION:
			version_str = optarg;
			version = strtoul(optarg, &e, 0);
			if (!*optarg || (e && *e)) {
				fprintf(stderr, "Invalid --version\n");
				parse_error = 1;
			}
			break;

		case OPT_MINVERSION:
			min_version = strtoul(optarg, &e, 0);
			if (!*optarg || (e && *e)) {
				fprintf(stderr, "Invalid --minversion\n");
				parse_error = 1;
			}
			break;

		case OPT_PAD:
			opt_pad = strtoul(optarg, &e, 0);
			if (!*optarg || (e && *e)) {
				fprintf(stderr, "Invalid --pad\n");
				parse_error = 1;
			}
			break;
		}
	}

	if (parse_error) {
		print_help(argv[0]);
		return 1;
	}

	switch (mode) {
	case OPT_MODE_PACK:

		if (!keyblock_file)
			Fatal("Missing required keyblock file.\n");

		t_keyblock = (VbKeyBlockHeader *)ReadFile(keyblock_file, 0);
		if (!t_keyblock)
			Fatal("Error reading key block.\n");

		if (!signprivkey_file)
			Fatal("Missing required signprivate file.\n");

		signpriv_key = PrivateKeyRead(signprivkey_file);
		if (!signpriv_key)
			Fatal("Error reading signing key.\n");

		if (!config_file)
			Fatal("Missing required config file.\n");

		Debug("Reading %s\n", config_file);
		t_config_data =
			ReadConfigFile(config_file, &t_config_size);
		if (!t_config_data)
			Fatal("Error reading config file.\n");

		if (!bootloader_file)
			Fatal("Missing required bootloader file.\n");

		Debug("Reading %s\n", bootloader_file);
		t_bootloader_data = ReadFile(bootloader_file,
					     &t_bootloader_size);
		if (!t_bootloader_data)
			Fatal("Error reading bootloader file.\n");
		Debug(" bootloader file size=0x%" PRIx64 "\n",
		      t_bootloader_size);

		if (!vmlinuz_file)
			Fatal("Missing required vmlinuz file.\n");
		Debug("Reading %s\n", vmlinuz_file);
		vmlinuz_buf = ReadFile(vmlinuz_file, &vmlinuz_size);
		if (!vmlinuz_buf)
			Fatal("Error reading vmlinuz file.\n");
		Debug(" vmlinuz file size=0x%" PRIx64 "\n",
		      vmlinuz_size);
		if (!vmlinuz_size)
			Fatal("Empty vmlinuz file\n");

		kblob_data = CreateKernelBlob(
			vmlinuz_buf, vmlinuz_size,
			arch, kernel_body_load_address,
			t_config_data, t_config_size,
			t_bootloader_data, t_bootloader_size,
			&kblob_size);
		if (!kblob_data)
			Fatal("Unable to create kernel blob\n");

		Debug("kblob_size = 0x%" PRIx64 "\n", kblob_size);

		vblock_data = SignKernelBlob(kblob_data, kblob_size, opt_pad,
					     version, kernel_body_load_address,
					     t_keyblock, signpriv_key,
					     &vblock_size);
		if (!vblock_data)
			Fatal("Unable to sign kernel blob\n");

		Debug("vblock_size = 0x%" PRIx64 "\n", vblock_size);

		if (opt_vblockonly)
			rv = WriteSomeParts(filename,
					    vblock_data, vblock_size,
					    NULL, 0);
		else
			rv = WriteSomeParts(filename,
					    vblock_data, vblock_size,
					    kblob_data, kblob_size);
		return rv;

	case OPT_MODE_REPACK:

		/* Required */

		if (!signprivkey_file)
			Fatal("Missing required signprivate file.\n");

		signpriv_key = PrivateKeyRead(signprivkey_file);
		if (!signpriv_key)
			Fatal("Error reading signing key.\n");

		if (!oldfile)
			Fatal("Missing previously packed blob.\n");

		/* Load the kernel partition */
		kpart_data = ReadOldKPartFromFileOrDie(oldfile, &kpart_size);

		kblob_data = UnpackKPart(kpart_data, kpart_size, opt_pad,
					 &keyblock, &preamble, &kblob_size);

		if (!kblob_data)
			Fatal("Unable to unpack kernel partition\n");

		kernel_body_load_address = preamble->body_load_address;

		/* Update the config if asked */
		if (config_file) {
			Debug("Reading %s\n", config_file);
			t_config_data =
				ReadConfigFile(config_file, &t_config_size);
			if (!t_config_data)
				Fatal("Error reading config file.\n");
			if (0 != UpdateKernelBlobConfig(
				    kblob_data, kblob_size,
				    t_config_data, t_config_size))
				Fatal("Unable to update config\n");
		}

		if (!version_str)
			version = preamble->kernel_version;

		if (keyblock_file) {
			t_keyblock =
				(VbKeyBlockHeader *)ReadFile(keyblock_file, 0);
			if (!t_keyblock)
				Fatal("Error reading key block.\n");
		}

		/* Reuse previous body size */
		vblock_data = SignKernelBlob(kblob_data, kblob_size, opt_pad,
					     version, kernel_body_load_address,
					     t_keyblock ? t_keyblock : keyblock,
					     signpriv_key, &vblock_size);
		if (!vblock_data)
			Fatal("Unable to sign kernel blob\n");

		if (opt_vblockonly)
			rv = WriteSomeParts(filename,
					    vblock_data, vblock_size,
					    NULL, 0);
		else
			rv = WriteSomeParts(filename,
					    vblock_data, vblock_size,
					    kblob_data, kblob_size);
		return rv;

	case OPT_MODE_VERIFY:

		/* Optional */

		if (signpubkey_file) {
			signpub_key = PublicKeyRead(signpubkey_file);
			if (!signpub_key)
				Fatal("Error reading public key.\n");
		}

		/* Do it */

		/* Load the kernel partition */
		kpart_data = ReadOldKPartFromFileOrDie(filename, &kpart_size);

		kblob_data = UnpackKPart(kpart_data, kpart_size, opt_pad,
					 0, 0, &kblob_size);
		if (!kblob_data)
			Fatal("Unable to unpack kernel partition\n");

		rv = VerifyKernelBlob(kblob_data, kblob_size,
				      signpub_key, keyblock_file, min_version);

		return rv;
	}

	fprintf(stderr,
		"You must specify a mode: --pack, --repack or --verify\n");
	print_help(argv[0]);
	return 1;
}

DECLARE_FUTIL_COMMAND(vbutil_kernel2, do_vbutil_kernel2,
		      "Creates, signs, and verifies the kernel partition",
		      print_help);