/*
 * Copyright (c) 2007, Google Inc.
 *
 * Copyright (c) 2013, 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;
}