/*-
 * SPDX-License-Identifier: BSD-2-Clause
 *
 * Copyright (c) 2013,2014 Juniper Networks, Inc.
 * 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.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``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 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 <sys/param.h>
#include <sys/stat.h>
#include <errno.h>
#include <err.h>
#include <fcntl.h>
#include <getopt.h>
#include <libutil.h>
#include <limits.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sysexits.h>
#include <unistd.h>

#include "image.h"
#include "format.h"
#include "mkimg.h"
#include "scheme.h"

#define LONGOPT_FORMATS         0x01000001
#define LONGOPT_SCHEMES         0x01000002
#define LONGOPT_VERSION         0x01000003
#define LONGOPT_CAPACITY        0x01000004

static struct option longopts[] = {
        { "formats", no_argument, NULL, LONGOPT_FORMATS },
        { "schemes", no_argument, NULL, LONGOPT_SCHEMES },
        { "version", no_argument, NULL, LONGOPT_VERSION },
        { "capacity", required_argument, NULL, LONGOPT_CAPACITY },
        { NULL, 0, NULL, 0 }
};

static uint64_t min_capacity = 0;
static uint64_t max_capacity = 0;

/* Fixed timestamp for reproducible builds. */
time_t timestamp = (time_t)-1;

struct partlisthead partlist = TAILQ_HEAD_INITIALIZER(partlist);
u_int nparts = 0;

u_int unit_testing;
u_int verbose;

u_int ncyls = 0;
u_int nheads = 1;
u_int nsecs = 1;
u_int secsz = 512;
u_int blksz = 0;
uint32_t active_partition = 0;

static void
print_formats(int usage)
{
        struct mkimg_format *f;
        const char *sep;

        if (usage) {
                fprintf(stderr, "    formats:\n");
                f = NULL;
                while ((f = format_iterate(f)) != NULL) {
                        fprintf(stderr, "\t%s\t-  %s\n", f->name,
                            f->description);
                }
        } else {
                sep = "";
                f = NULL;
                while ((f = format_iterate(f)) != NULL) {
                        printf("%s%s", sep, f->name);
                        sep = " ";
                }
                putchar('\n');
        }
}

static void
print_schemes(int usage)
{
        struct mkimg_scheme *s;
        const char *sep;

        if (usage) {
                fprintf(stderr, "    schemes:\n");
                s = NULL;
                while ((s = scheme_iterate(s)) != NULL) {
                        fprintf(stderr, "\t%s\t-  %s\n", s->name,
                            s->description);
                }
        } else {
                sep = "";
                s = NULL;
                while ((s = scheme_iterate(s)) != NULL) {
                        printf("%s%s", sep, s->name);
                        sep = " ";
                }
                putchar('\n');
        }
}

static void
print_version(void)
{
        u_int width;

#ifdef __LP64__
        width = 64;
#else
        width = 32;
#endif
        printf("mkimg %u (%u-bit)\n", MKIMG_VERSION, width);
}

static void
usage(const char *why)
{

        if (why != NULL) {
                warnx("error: %s", why);
                fputc('\n', stderr);
        }
        fprintf(stderr, "usage: %s <options>\n", getprogname());

        fprintf(stderr, "    options:\n");
        fprintf(stderr, "\t--formats\t-  list image formats\n");
        fprintf(stderr, "\t--schemes\t-  list partition schemes\n");
        fprintf(stderr, "\t--version\t-  show version information\n");
        fprintf(stderr, "\t--capacity\t-  minimum and maximum capacity (in bytes)\n");
        fputc('\n', stderr);
        fprintf(stderr, "\t-a <num>\t-  mark num'th partition as active\n");
        fprintf(stderr, "\t-b <file>\t-  file containing boot code\n");
        fprintf(stderr, "\t-c <num>\t-  minimum capacity (in bytes) of the disk\n");
        fprintf(stderr, "\t-C <num>\t-  maximum capacity (in bytes) of the disk\n");
        fprintf(stderr, "\t-f <format>\n");
        fprintf(stderr, "\t-h\t\t-  show this usage information\n");
        fprintf(stderr, "\t-o <file>\t-  file to write image into\n");
        fprintf(stderr, "\t-p <partition>\n");
        fprintf(stderr, "\t-s <scheme>\n");
        fprintf(stderr, "\t-t <num>\t-  set timestamp (seconds since epoch)\n");
        fprintf(stderr, "\t-v\t\t-  increase verbosity\n");
        fprintf(stderr, "\t-y\t\t-  [developers] enable unit test\n");
        fprintf(stderr, "\t-H <num>\t-  number of heads to simulate\n");
        fprintf(stderr, "\t-P <num>\t-  physical sector size\n");
        fprintf(stderr, "\t-S <num>\t-  logical sector size\n");
        fprintf(stderr, "\t-T <num>\t-  number of tracks to simulate\n");
        fputc('\n', stderr);
        print_formats(1);
        fputc('\n', stderr);
        print_schemes(1);
        fputc('\n', stderr);
        fprintf(stderr, "    partition specification:\n");
        fprintf(stderr, "\t<type>[/<label>]::<size>[:[+]<offset>]\t-  "
            "empty partition of given size and\n\t\t\t\t\t"
            "   optional relative or absolute offset\n");
        fprintf(stderr, "\t<type>[/<label>]:=<file>[:[+]offset]\t-  partition "
            "content and size are\n\t\t\t\t\t"
            "   determined by the named file and\n"
            "\t\t\t\t\t   optional relative or absolute offset\n");
        fprintf(stderr, "\t<type>[/<label>]:-<cmd>\t\t-  partition content and size "
            "are taken\n\t\t\t\t\t   from the output of the command to run\n");
        fprintf(stderr, "\t-\t\t\t\t-  unused partition entry\n");
        fprintf(stderr, "\t    where:\n");
        fprintf(stderr, "\t\t<type>\t-  scheme neutral partition type\n");
        fprintf(stderr, "\t\t<label>\t-  optional scheme-dependent partition "
            "label\n");

        exit(EX_USAGE);
}

static int
parse_uint32(uint32_t *valp, uint32_t min, uint32_t max, const char *arg)
{
        uint64_t val;

        if (expand_number(arg, &val) == -1)
                return (errno);
        if (val > UINT_MAX || val < (uint64_t)min || val > (uint64_t)max)
                return (EINVAL);
        *valp = (uint32_t)val;
        return (0);
}

static int
parse_uint64(uint64_t *valp, uint64_t min, uint64_t max, const char *arg)
{
        uint64_t val;

        if (expand_number(arg, &val) == -1)
                return (errno);
        if (val < min || val > max)
                return (EINVAL);
        *valp = val;
        return (0);
}

static int
pwr_of_two(u_int nr)
{

        return (((nr & (nr - 1)) == 0) ? 1 : 0);
}

/*
 * A partition specification has the following format:
 *      <type> ':' <kind> <contents>
 * where:
 *      type      the partition type alias
 *      kind      the interpretation of the contents specification
 *                ':'   contents holds the size of an empty partition
 *                '='   contents holds the name of a file to read
 *                '-'   contents holds a command to run; the output of
 *                      which is the contents of the partition.
 *      contents  the specification of a partition's contents
 *
 * A specification that is a single dash indicates an unused partition
 * entry.
 */
static int
parse_part(const char *spec)
{
        struct part *part;
        char *sep;
        size_t len;
        int error;

        if (strcmp(spec, "-") == 0) {
                nparts++;
                return (0);
        }

        part = calloc(1, sizeof(struct part));
        if (part == NULL)
                return (ENOMEM);

        sep = strchr(spec, ':');
        if (sep == NULL) {
                error = EINVAL;
                goto errout;
        }
        len = sep - spec + 1;
        if (len < 2) {
                error = EINVAL;
                goto errout;
        }
        part->alias = malloc(len);
        if (part->alias == NULL) {
                error = ENOMEM;
                goto errout;
        }
        strlcpy(part->alias, spec, len);
        spec = sep + 1;

        switch (*spec) {
        case ':':
                part->kind = PART_KIND_SIZE;
                break;
        case '=':
                part->kind = PART_KIND_FILE;
                break;
        case '-':
                part->kind = PART_KIND_PIPE;
                break;
        default:
                error = EINVAL;
                goto errout;
        }
        spec++;

        part->contents = strdup(spec);
        if (part->contents == NULL) {
                error = ENOMEM;
                goto errout;
        }

        spec = part->alias;
        sep = strchr(spec, '/');
        if (sep != NULL) {
                *sep++ = '\0';
                if (strlen(part->alias) == 0 || strlen(sep) == 0) {
                        error = EINVAL;
                        goto errout;
                }
                part->label = strdup(sep);
                if (part->label == NULL) {
                        error = ENOMEM;
                        goto errout;
                }
        }

        part->index = nparts;
        TAILQ_INSERT_TAIL(&partlist, part, link);
        nparts++;
        return (0);

 errout:
        if (part->alias != NULL)
                free(part->alias);
        free(part);
        return (error);
}

#if defined(SPARSE_WRITE)
ssize_t
sparse_write(int fd, const void *ptr, size_t sz)
{
        const char *buf, *p;
        off_t ofs;
        size_t len;
        ssize_t wr, wrsz;

        buf = ptr;
        wrsz = 0;
        p = memchr(buf, 0, sz);
        while (sz > 0) {
                len = (p != NULL) ? (size_t)(p - buf) : sz;
                if (len > 0) {
                        len = (len + secsz - 1) & ~(secsz - 1);
                        if (len > sz)
                                len = sz;
                        wr = write(fd, buf, len);
                        if (wr < 0)
                                return (-1);
                } else {
                        while (len < sz && *p++ == '\0')
                                len++;
                        if (len < sz)
                                len &= ~(secsz - 1);
                        if (len == 0)
                                continue;
                        ofs = lseek(fd, len, SEEK_CUR);
                        if (ofs < 0)
                                return (-1);
                        wr = len;
                }
                buf += wr;
                sz -= wr;
                wrsz += wr;
                p = memchr(buf, 0, sz);
        }
        return (wrsz);
}
#endif /* SPARSE_WRITE */

void
mkimg_chs(lba_t lba, u_int maxcyl, u_int *cylp, u_int *hdp, u_int *secp)
{
        u_int hd, sec;

        *cylp = *hdp = *secp = ~0U;
        if (nsecs == 1 || nheads == 1)
                return;

        sec = lba % nsecs + 1;
        lba /= nsecs;
        hd = lba % nheads;
        lba /= nheads;
        if (lba > maxcyl)
                return;

        *cylp = lba;
        *hdp = hd;
        *secp = sec;
}

static int
capacity_resize(lba_t end)
{
        lba_t min_capsz, max_capsz;

        min_capsz = (min_capacity + secsz - 1) / secsz;
        max_capsz = (max_capacity + secsz - 1) / secsz;

        if (max_capsz != 0 && end > max_capsz)
                return (ENOSPC);
        if (end >= min_capsz)
                return (0);

        return (image_set_size(min_capsz));
}

static void
mkimg_validate(void)
{
        struct part *part, *part2;
        lba_t start, end, start2, end2;
        int i, j;

        i = 0;

        TAILQ_FOREACH(part, &partlist, link) {
                start = part->block;
                end = part->block + part->size;
                j = i + 1;
                part2 = TAILQ_NEXT(part, link);
                if (part2 == NULL)
                        break;

                TAILQ_FOREACH_FROM(part2, &partlist, link) {
                        start2 = part2->block;
                        end2 = part2->block + part2->size;

                        if ((start >= start2 && start < end2) ||
                            (end > start2 && end <= end2)) {
                                errx(1, "partition %d overlaps partition %d",
                                    i, j);
                        }

                        j++;
                }

                i++;
        }
}

static void
mkimg(void)
{
        FILE *fp;
        struct part *part;
        lba_t block, blkoffset;
        uint64_t bytesize, byteoffset;
        char *size, *offset;
        bool abs_offset;
        int error, fd;

        /* First check partition information */
        TAILQ_FOREACH(part, &partlist, link) {
                error = scheme_check_part(part);
                if (error)
                        errc(EX_DATAERR, error, "partition %d", part->index+1);
        }

        block = scheme_metadata(SCHEME_META_IMG_START, 0);
        abs_offset = false;
        TAILQ_FOREACH(part, &partlist, link) {
                byteoffset = blkoffset = 0;
                abs_offset = false;

                /* Look for an offset. Set size too if we can. */
                switch (part->kind) {
                case PART_KIND_SIZE:
                case PART_KIND_FILE:
                        offset = part->contents;
                        size = strsep(&offset, ":");
                        if (part->kind == PART_KIND_SIZE &&
                            expand_number(size, &bytesize) == -1)
                                error = errno;
                        if (offset != NULL) {
                                if (*offset != '+')
                                        abs_offset = true;
                                else
                                        offset++;
                                if (expand_number(offset, &byteoffset) == -1)
                                        error = errno;
                        }
                        break;
                }

                /* Work out exactly where the partition starts. */
                blkoffset = (byteoffset + secsz - 1) / secsz;
                if (abs_offset)
                        block = scheme_metadata(SCHEME_META_PART_ABSOLUTE,
                            blkoffset);
                else
                        block = scheme_metadata(SCHEME_META_PART_BEFORE,
                            block + blkoffset);
                part->block = block;

                if (verbose)
                        fprintf(stderr, "partition %d: starting block %llu "
                            "... ", part->index + 1, (long long)part->block);

                /* Pull in partition contents, set size if we haven't yet. */
                switch (part->kind) {
                case PART_KIND_FILE:
                        fd = open(part->contents, O_RDONLY, 0);
                        if (fd != -1) {
                                error = image_copyin(block, fd, &bytesize);
                                close(fd);
                        } else
                                error = errno;
                        break;
                case PART_KIND_PIPE:
                        fp = popen(part->contents, "r");
                        if (fp != NULL) {
                                fd = fileno(fp);
                                error = image_copyin(block, fd, &bytesize);
                                pclose(fp);
                        } else
                                error = errno;
                        break;
                }
                if (error)
                        errc(EX_IOERR, error, "partition %d", part->index + 1);
                part->size = (bytesize + secsz - 1) / secsz;
                if (verbose) {
                        bytesize = part->size * secsz;
                        fprintf(stderr, "size %llu bytes (%llu blocks)\n",
                             (long long)bytesize, (long long)part->size);
                        if (abs_offset) {
                                fprintf(stderr,
                                    "    location %llu bytes (%llu blocks)\n",
                                    (long long)byteoffset,
                                    (long long)blkoffset);
                        } else if (blkoffset > 0) {
                                fprintf(stderr,
                                    "    offset %llu bytes (%llu blocks)\n",
                                    (long long)byteoffset,
                                    (long long)blkoffset);
                        }
                }
                block = scheme_metadata(SCHEME_META_PART_AFTER,
                    part->block + part->size);
        }

        mkimg_validate();

        block = scheme_metadata(SCHEME_META_IMG_END, block);
        error = image_set_size(block);
        if (!error) {
                error = capacity_resize(block);
                block = image_get_size();
        }
        if (!error) {
                error = format_resize(block);
                block = image_get_size();
        }
        if (error)
                errc(EX_IOERR, error, "image sizing");
        ncyls = block / (nsecs * nheads);
        error = scheme_write(block);
        if (error)
                errc(EX_IOERR, error, "writing metadata");
}

int
main(int argc, char *argv[])
{
        const char *format_name;
        int bcfd, outfd;
        int c, error;

        bcfd = -1;
        outfd = 1;      /* Write to stdout by default */
        while ((c = getopt_long(argc, argv, "a:b:c:C:f:ho:p:s:t:vyH:P:S:T:",
            longopts, NULL)) != -1) {
                switch (c) {
                case 'a':       /* ACTIVE PARTITION, if supported */
                        error = parse_uint32(&active_partition, 1, 100, optarg);
                        if (error)
                                errc(EX_DATAERR, error, "Partition ordinal");
                        break;
                case 'b':       /* BOOT CODE */
                        if (bcfd != -1)
                                usage("multiple bootcode given");
                        bcfd = open(optarg, O_RDONLY, 0);
                        if (bcfd == -1)
                                err(EX_UNAVAILABLE, "%s", optarg);
                        break;
                case 'c':       /* MINIMUM CAPACITY */
                        error = parse_uint64(&min_capacity, 1, INT64_MAX, optarg);
                        if (error)
                                errc(EX_DATAERR, error, "minimum capacity in bytes");
                        break;
                case 'C':       /* MAXIMUM CAPACITY */
                        error = parse_uint64(&max_capacity, 1, INT64_MAX, optarg);
                        if (error)
                                errc(EX_DATAERR, error, "maximum capacity in bytes");
                        break;
                case 'f':       /* OUTPUT FORMAT */
                        if (format_selected() != NULL)
                                usage("multiple formats given");
                        error = format_select(optarg);
                        if (error)
                                errc(EX_DATAERR, error, "format");
                        break;
                case 'h':       /* HELP */
                        usage(NULL);
                        break;
                case 'o':       /* OUTPUT FILE */
                        if (outfd != 1)
                                usage("multiple output files given");
                        outfd = open(optarg, O_WRONLY | O_CREAT | O_TRUNC,
                            S_IWUSR | S_IRUSR | S_IRGRP | S_IROTH);
                        if (outfd == -1)
                                err(EX_CANTCREAT, "%s", optarg);
                        break;
                case 'p':       /* PARTITION */
                        error = parse_part(optarg);
                        if (error)
                                errc(EX_DATAERR, error, "partition");
                        break;
                case 's':       /* SCHEME */
                        if (scheme_selected() != NULL)
                                usage("multiple schemes given");
                        error = scheme_select(optarg);
                        if (error)
                                errc(EX_DATAERR, error, "scheme");
                        break;
                case 't': {
                        char *ep;
                        long long val;

                        errno = 0;
                        val = strtoll(optarg, &ep, 0);
                        if (ep == optarg || *ep != '\0')
                                errno = EINVAL;
                        if (errno != 0)
                                errc(EX_DATAERR, errno, "timestamp");
                        timestamp = (time_t)val;
                        break;
                }
                case 'y':
                        unit_testing++;
                        break;
                case 'v':
                        verbose++;
                        break;
                case 'H':       /* GEOMETRY: HEADS */
                        error = parse_uint32(&nheads, 1, 255, optarg);
                        if (error)
                                errc(EX_DATAERR, error, "number of heads");
                        break;
                case 'P':       /* GEOMETRY: PHYSICAL SECTOR SIZE */
                        error = parse_uint32(&blksz, 512, INT_MAX+1U, optarg);
                        if (error == 0 && !pwr_of_two(blksz))
                                error = EINVAL;
                        if (error)
                                errc(EX_DATAERR, error, "physical sector size");
                        break;
                case 'S':       /* GEOMETRY: LOGICAL SECTOR SIZE */
                        error = parse_uint32(&secsz, 512, INT_MAX+1U, optarg);
                        if (error == 0 && !pwr_of_two(secsz))
                                error = EINVAL;
                        if (error)
                                errc(EX_DATAERR, error, "logical sector size");
                        break;
                case 'T':       /* GEOMETRY: TRACK SIZE */
                        error = parse_uint32(&nsecs, 1, 63, optarg);
                        if (error)
                                errc(EX_DATAERR, error, "track size");
                        break;
                case LONGOPT_FORMATS:
                        print_formats(0);
                        exit(EX_OK);
                        /*NOTREACHED*/
                case LONGOPT_SCHEMES:
                        print_schemes(0);
                        exit(EX_OK);
                        /*NOTREACHED*/
                case LONGOPT_VERSION:
                        print_version();
                        exit(EX_OK);
                        /*NOTREACHED*/
                case LONGOPT_CAPACITY:
                        error = parse_uint64(&min_capacity, 1, INT64_MAX, optarg);
                        if (error)
                                errc(EX_DATAERR, error, "capacity in bytes");
                        max_capacity = min_capacity;
                        break;
                default:
                        usage("unknown option");
                }
        }

        if (argc > optind)
                usage("trailing arguments");
        if (scheme_selected() == NULL && nparts > 0)
                usage("no scheme");
        if (nparts == 0 && min_capacity == 0)
                usage("no partitions");
        if (max_capacity != 0 && min_capacity > max_capacity)
                usage("minimum capacity cannot be larger than the maximum one");

        if (secsz > blksz) {
                if (blksz != 0)
                        errx(EX_DATAERR, "the physical block size cannot "
                            "be smaller than the sector size");
                blksz = secsz;
        }

        if (secsz > scheme_max_secsz())
                errx(EX_DATAERR, "maximum sector size supported is %u; "
                    "size specified is %u", scheme_max_secsz(), secsz);

        if (nparts > scheme_max_parts())
                errx(EX_DATAERR, "%d partitions supported; %d given",
                    scheme_max_parts(), nparts);

        if (format_selected() == NULL)
                format_select("raw");

        if (bcfd != -1) {
                error = scheme_bootcode(bcfd);
                close(bcfd);
                if (error)
                        errc(EX_DATAERR, error, "boot code");
        }

        format_name = format_selected()->name;
        if (verbose) {
                fprintf(stderr, "Logical sector size: %u\n", secsz);
                fprintf(stderr, "Physical block size: %u\n", blksz);
                fprintf(stderr, "Sectors per track:   %u\n", nsecs);
                fprintf(stderr, "Number of heads:     %u\n", nheads);
                fputc('\n', stderr);
                if (scheme_selected())
                        fprintf(stderr, "Partitioning scheme: %s\n",
                            scheme_selected()->name);
                fprintf(stderr, "Output file format:  %s\n",
                    format_name);
                fputc('\n', stderr);
        }

#if defined(SPARSE_WRITE)
        /*
         * sparse_write() fails if output is not seekable so fail early
         * not wasting some load unless output format is raw
         */
        if (strcmp("raw", format_name) &&
            lseek(outfd, (off_t)0, SEEK_CUR) == -1 && errno == ESPIPE)
                errx(EX_USAGE, "%s: output must be seekable", format_name);
#endif

        error = image_init();
        if (error)
                errc(EX_OSERR, error, "cannot initialize");

        mkimg();

        if (verbose) {
                fputc('\n', stderr);
                fprintf(stderr, "Number of cylinders: %u\n", ncyls);
        }

        error = format_write(outfd);
        if (error)
                errc(EX_IOERR, error, "writing image");

        return (0);
}