/*
 * Copyright (c) 2007, Google Inc.
 * All rights reserved.
 *
 * Copyright (c) 2009-2011, The Linux Foundation. All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are
 * met:
 *    * Redistributions of source code must retain the above copyright
 *      notice, this list of conditions and the following disclaimer.
 *    * 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.
 *    * Neither the name of The Linux Foundation nor the names of its
 *      contributors may be used to endorse or promote products derived
 *      from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED "AS IS" AND ANY EXPRESS OR IMPLIED
 * WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NON-INFRINGEMENT
 * ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS
 * 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 <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <fcntl.h>
#include <string.h>

#include <sys/stat.h>

int print_usage()
{
	fprintf(stderr, "usage: mkheader <bin> <hdr> <none|unified-boot>\n");
	fprintf(stderr, "       mkheader <bin> <hdr> <unsecure-boot>"
		" <outbin>\n");
	fprintf(stderr, "       mkheader <bin> <hdr> <secure-boot> <outbin>"
		" <maxsize>\n");
	fprintf(stderr, "       mkheader <bin> <hdr> <secure-boot> <outbin>"
		" <maxsize> <certchain> <files...>\n\n");
	fprintf(stderr, "bin:               Input raw appsbl binary\n");
	fprintf(stderr,
		"hdr:               Output of appsbl header location\n");
	fprintf(stderr,
		"outbin:            Output of the signed or unsigned"
		" apps boot location\n");
	fprintf(stderr,
		"maxsize:           Maximum size for certificate" " chain\n");
	fprintf(stderr,
		"certchain:         Output of the certchain location\n");
	fprintf(stderr,
		"files:             Input format <bin signature>"
		" <certifcate file(s) for certificate chain>...\n");
	fprintf(stderr,
		"certificate chain: Files will be concatenated in order"
		" to create the certificate chain\n\n");
	return -1;
}

int cat(FILE * in, FILE * out, unsigned size, unsigned buff_size)
{
	unsigned bytes_left = size;
	char buf[buff_size];
	int ret = 0;

	while (bytes_left) {
		fread(buf, sizeof(char), buff_size, in);
		if (!feof(in)) {
			bytes_left -= fwrite(buf, sizeof(char), buff_size, out);
		} else
			bytes_left = 0;
	}
	ret = ferror(in) | ferror(out);
	if (ret)
		fprintf(stderr, "ERROR: Occured during file concatenation\n");
	return ret;
}

int main(int argc, char *argv[])
{
	struct stat s;
	unsigned size, base;
	int unified_boot = 0;
	unsigned unified_boot_magic[20];
	unsigned non_unified_boot_magic[10];
	unsigned magic_len = 0;
	unsigned *magic;
	unsigned cert_chain_size = 0;
	unsigned signature_size = 0;
	int secure_boot = 0;
	int fd;

	if (argc < 3) {
		return print_usage();
	}

	if (argc == 4) {
		if (!strcmp("unified-boot", argv[3])) {
			unified_boot = 1;
		} else if (!strcmp("secure-boot", argv[3])) {
			fprintf(stderr,
				"ERROR: Missing arguments: [outbin maxsize] |"
				" [outbin, maxsize, certchain,"
				" signature + certifcate(s)]\n");
			return print_usage();
		} else if (!strcmp("unsecure-boot", argv[3])) {
			fprintf(stderr, "ERROR: Missing arguments:"
				" outbin directory\n");
			return print_usage();
		}
	}

	if (argc > 4) {
		if (!strcmp("secure-boot", argv[3])) {
			if (argc < 9 && argc != 6) {
				fprintf(stderr,
					"ERROR: Missing argument(s):"
					" [outbin maxsize] | [outbin, maxsize,"
					" certchain,"
					" signature + certifcate(s)]\n");
				return print_usage();
			}
			secure_boot = 1;
			signature_size = 256;	//Support SHA 256
			cert_chain_size = atoi(argv[5]);
		}
	}

	if (stat(argv[1], &s)) {
		perror("cannot stat binary");
		return -1;
	}

	if (unified_boot) {
		magic = unified_boot_magic;
		magic_len = sizeof(unified_boot_magic);
	} else {
		magic = non_unified_boot_magic;
		magic_len = sizeof(non_unified_boot_magic);
	}

	size = s.st_size;
#if MEMBASE
	base = MEMBASE;
#else
	base = 0;
#endif

	printf("Image Destination Pointer: 0x%x\n", base);

	magic[0] = 0x00000005;	/* appsbl */
	magic[1] = 0x00000003;	//Flash_partition_version /* nand */
	magic[2] = 0x00000000;	//image source pointer
	magic[3] = base;	//image destination pointer
	magic[4] = size + cert_chain_size + signature_size;	//image size
	magic[5] = size;	//code size
	magic[6] = base + size;
	magic[7] = signature_size;
	magic[8] = size + base + signature_size;
	magic[9] = cert_chain_size;

	if (unified_boot == 1) {
		magic[10] = 0x33836685;	/* cookie magic number */
		magic[11] = 0x00000001;	/* cookie version */
		magic[12] = 0x00000002;	/* file formats */
		magic[13] = 0x00000000;
		magic[14] = 0x00000000;	/* not setting size for boot.img */
		magic[15] = 0x00000000;
		magic[16] = 0x00000000;
		magic[17] = 0x00000000;
		magic[18] = 0x00000000;
		magic[19] = 0x00000000;
	}

	fd = open(argv[2], O_WRONLY | O_CREAT | O_TRUNC, 0644);
	if (fd < 0) {
		perror("cannot open header for writing");
		return -1;
	}
	if (write(fd, magic, magic_len) != magic_len) {
		perror("cannot write header");
		close(fd);
		unlink(argv[2]);
		return -1;
	}
	close(fd);

	if (secure_boot && argc > 6) {
		FILE *input_file;
		FILE *output_file;
		unsigned buff_size = 1;
		char buf[buff_size];
		unsigned bytes_left;
		unsigned current_cert_chain_size = 0;
		int padding_size = 0;
		int i;

		if ((output_file = fopen(argv[6], "wb")) == NULL) {
			perror("ERROR: Occured during fopen");
			return -1;
		}
		printf("Certificate Chain Output File: %s\n", argv[6]);

		for (i = 8; i < argc; i++) {
			if ((input_file = fopen(argv[i], "rb")) == NULL) {
				perror("ERROR: Occured during fopen");
				return -1;
			}
			stat(argv[i], &s);
			bytes_left = s.st_size;
			current_cert_chain_size += bytes_left;
			if (cat(input_file, output_file, bytes_left, buff_size))
				return -1;
			fclose(input_file);
		}

		//Pad certifcate chain to the max expected size from input
		memset(buf, 0xFF, sizeof(buf));
		padding_size = cert_chain_size - current_cert_chain_size;

		if (padding_size < 0) {
			fprintf(stderr, "ERROR: Input certificate chain"
				" (Size=%d) is larger than the maximum"
				" specified (Size=%d)\n",
				current_cert_chain_size, cert_chain_size);
			return -1;
		}

		bytes_left = (padding_size > 0) ? padding_size : 0;
		while (bytes_left) {
			if (!ferror(output_file))
				bytes_left -= fwrite(buf,
						     sizeof(buf),
						     buff_size, output_file);
			else {
				fprintf(stderr, "ERROR: Occured during"
					" certifcate chain padding\n");
				return -1;
			}
		}
		fclose(output_file);

		/* Concat and combine to signed image.
		 * Format [HDR][RAW APPSBOOT][PADDED CERT CHAIN]
		 */
		if ((output_file = fopen(argv[4], "wb")) == NULL) {
			perror("ERROR: Occured during fopen");
			return -1;
		}
		printf("Image Output File: %s\n", argv[4]);

		//Header
		if ((input_file = fopen(argv[2], "rb")) == NULL) {
			perror("ERROR: Occured during fopen");
			return -1;
		}
		stat(argv[2], &s);
		if (cat(input_file, output_file, s.st_size, buff_size))
			return -1;
		fclose(input_file);

		//Raw Appsbl
		if ((input_file = fopen(argv[1], "rb")) == NULL) {
			perror("ERROR: Occured during fopen");
			return -1;
		}
		stat(argv[1], &s);
		if (cat(input_file, output_file, s.st_size, buff_size))
			return -1;
		fclose(input_file);

		//Signature
		if ((input_file = fopen(argv[7], "rb")) == NULL) {
			perror("ERROR: Occured during fopen");
			return -1;
		}
		stat(argv[7], &s);
		if (cat(input_file, output_file, s.st_size, buff_size))
			return -1;
		fclose(input_file);

		//Certifcate Chain
		if ((input_file = fopen(argv[6], "rb")) == NULL) {
			perror("ERROR: Occured during fopen");
			return -1;
		}
		if (cat(input_file, output_file,
			(current_cert_chain_size + padding_size), buff_size))
			return -1;
		fclose(input_file);

		fclose(output_file);

	} else if (argc == 5 || argc == 6) {
		FILE *input_file;
		FILE *output_file;
		unsigned buff_size = 1;
		char buf[buff_size];

		/* Concat and combine to unsigned image.
		 * Format [HDR][RAW APPSBOOT]
		 */
		if ((output_file = fopen(argv[4], "wb")) == NULL) {
			perror("ERROR: Occured during fopen");
			return -1;
		}
		printf("Image Output File: %s\n", argv[4]);

		//Header
		if ((input_file = fopen(argv[2], "rb")) == NULL) {
			perror("ERROR: Occured during fopen");
			return -1;
		}
		stat(argv[2], &s);
		if (cat(input_file, output_file, s.st_size, buff_size))
			return -1;
		fclose(input_file);

		//Raw Appsbl
		if ((input_file = fopen(argv[1], "rb")) == NULL) {
			perror("ERROR: Occured during fopen");
			return -1;
		}
		stat(argv[1], &s);
		if (cat(input_file, output_file, s.st_size, buff_size))
			return -1;
		fclose(input_file);
		fclose(output_file);
	}

	printf("Done execution\n");

	return 0;
}