/*
 * Copyright (C) 2008 The Android Open Source Project
 * 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.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS 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
 * 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 <stdarg.h>
#include <string.h>
#include <errno.h>
#include <fcntl.h>
#include <unistd.h>
#include <limits.h>
#include <ctype.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <libgen.h>

#include "libzipfile/zipfile.h"
#include "fastboot.h"
#include "parser.h"

/* config file in system-converted tarball */
#define CONFIG_FILE  "manifest.txt"

/* config files in raw container generated by MIC */
#define FW_CONFIG    "ifwi.manifest"
#define PREOS_CONFIG "preos.manifest"
#define PLATFORM_IMG "platform.img"
#define DATA_IMG     "data.img"
#define CSA_IMG      "csa.img"

static usb_handle *usb = 0;
static const char *serial = 0;
static int wipe_data = 0;
static unsigned short vendor_id = 0;

void die(const char *fmt, ...)
{
    va_list ap;
    va_start(ap, fmt);
    fprintf(stderr,"error: ");
    vfprintf(stderr, fmt, ap);
    fprintf(stderr,"\n");
    va_end(ap);
    exit(1);
}

#ifdef _WIN32
void *load_file(const char *fn, unsigned *_sz);
#else
void *load_file(const char *fn, unsigned *_sz)
{
    char *data;
    int sz;
    int fd;
    int errno_tmp;

    data = 0;
    fd = open(fn, O_RDONLY);
    if(fd < 0) return 0;

    sz = lseek(fd, 0, SEEK_END);
    if(sz < 0) goto oops;

    if(lseek(fd, 0, SEEK_SET) != 0) goto oops;

    data = (char*) malloc(sz);
    if(data == 0) goto oops;

    if(read(fd, data, sz) != sz) goto oops;
    close(fd);

    if(_sz) *_sz = sz;
    return data;

oops:
    errno_tmp = errno;
    close(fd);
    if(data != 0) free(data);
    errno = errno_tmp;
    return 0;
}
#endif

int match_fastboot(usb_ifc_info *info)
{
    if(!(vendor_id && (info->dev_vendor == vendor_id)) &&
       (info->dev_vendor != 0x18d1) &&  // Google
       (info->dev_vendor != 0x0451) &&
       (info->dev_vendor != 0x0502) &&
       (info->dev_vendor != 0x0fce) &&  // Sony Ericsson
       (info->dev_vendor != 0x05c6) &&  // Qualcomm
       (info->dev_vendor != 0x22b8) &&  // Motorola
       (info->dev_vendor != 0x0955) &&  // Nvidia
       (info->dev_vendor != 0x413c) &&  // DELL
       (info->dev_vendor != 0x8087) &&  // Intel
       (info->dev_vendor != 0x04e8) &&  // Samsung
       (info->dev_vendor != 0x0bb4))    // HTC
            return -1;
    if(info->ifc_class != 0xff) return -1;
    if(info->ifc_subclass != 0x42) return -1;
    if(info->ifc_protocol != 0x03) return -1;
    // require matching serial number if a serial number is specified
    // at the command line with the -s option.
    if (serial && strcmp(serial, info->serial_number) != 0) return -1;
    return 0;
}

int list_devices_callback(usb_ifc_info *info)
{
    if (match_fastboot(info) == 0) {
        char* serial = info->serial_number;
        if (!info->writable) {
            serial = "no permissions"; // like "adb devices"
        }
        if (!serial[0]) {
            serial = "????????????";
        }
        // output compatible with "adb devices"
        printf("%s\tfastboot\n", serial);
    }

    return -1;
}

int check_usb_devices(usb_ifc_info *info);
usb_handle *open_device(void)
{
    static usb_handle *usb = 0;
    int announce = 1;

    if(usb) return usb;

    usb_open(check_usb_devices);

    for(;;) {
        usb = usb_open(match_fastboot);
        if(usb) return usb;
        if(announce) {
            announce = 0;
            fprintf(stderr,"< waiting for device >\n");
        }
        sleep(1);
    }
}

void list_devices(void) {
    // We don't actually open a USB device here,
    // just getting our callback called so we can
    // list all the connected devices.
    usb_open(list_devices_callback);
}

void usage(void)
{
    fprintf(stderr,
/*           1234567890123456789012345678901234567890123456789012345678901234567890123456 */
            "fastboot " VERSION "\n"
            "usage: fastboot [ <option> ] <command>\n"
            "\n"
            "commands:\n"
            "  [flashall] <filename>                    reflash device from a zip package\n"
            "  flash <partition> <filename>             write a file to a flash partition\n"
            "  erase <partition>                        erase a flash partition\n"
            "  getvar <variable>                        display a bootloader variable\n"
            "  oem <commands>                           oem specific commands\n"
            "  devices                                  list all connected devices\n"
            "  reboot                                   reboot device normally\n"
            "  reboot-bootloader                        reboot device into bootloader\n"
            "\n"
            "options:\n"
            "  -h|--help                                show this help message\n"
            "  -v|--version                             print fastboot version\n"
            "  -s|--serial <serial number>              specify device serial number\n"
            "  -i|--id <vendor id>                      specify a custom USB vendor id\n"
        );
}

int check_usb_devices(usb_ifc_info *info)
{
    static int nr_usb = 0;

    if (match_fastboot(info) == 0)
        nr_usb++;

    if (nr_usb > 1) {
        fprintf(stderr, "Error: found more than one devices connected.\n");
        list_devices();
        fprintf(stderr, "Please specify one by using '-s' option\n\n");
        usage();
        exit(EXIT_FAILURE);
    }

    return -1;
}

void *unzip_file(zipfile_t zip, const char *name, unsigned *sz)
{
    void *data;
    zipentry_t entry;
    unsigned datasz;

    entry = lookup_zipentry(zip, name);
    if (entry == NULL) {
        //fprintf(stderr, "archive does not contain '%s'\n", name);
        return 0;
    }

    *sz = get_zipentry_size(entry);

    datasz = *sz * 1.001;
    data = malloc(datasz);

    if(data == 0) {
        fprintf(stderr, "failed to allocate %d bytes\n", *sz);
        return 0;
    }

    if (decompress_zipentry(entry, data, datasz)) {
        fprintf(stderr, "failed to unzip '%s' from archive\n", name);
        free(data);
        return 0;
    }

    return data;
}

static char *strip(char *s)
{
    int n;
    while(*s && isspace(*s)) s++;
    n = strlen(s);
    while(n-- > 0) {
        if(!isspace(s[n])) break;
        s[n] = 0;
    }
    return s;
}

void queue_info_dump(void)
{
    fb_queue_notice("--------------------------------------------");
    fb_queue_display("preos", "Current Pre-OS Version ");
    fb_queue_display("ifwi",  "Current IFWI Version   ");
    fb_queue_notice("--------------------------------------------");
}

void do_flashall(char *fn)
{
    void *zdata;
    unsigned zsize;
    void *data;
    unsigned sz;
    zipfile_t zip;
    struct config conf;
    char ver[FB_RESPONSE_SZ + 1];
    char *pc;
    int status;

    /* get target IFWI major version */
    fprintf(stderr, "query system info...\n");
    fb_queue_query_save("ifwi", ver, sizeof(ver));
    usb = open_device();
    fb_execute_queue(usb);

    if ((pc = strchr(ver, '.')))
        *pc = 0;

    queue_info_dump();

    zdata = load_file(fn, &zsize);
    if (zdata == 0) die("failed to load '%s': %s", fn, strerror(errno));

    zip = init_zipfile(zdata, zsize);
    if(zip == 0) die("failed to access zipdata in '%s', is it a zip file?", fn);

    /* is converted-system tarball? */
    data = unzip_file(zip, CONFIG_FILE, &sz);
    if (data == 0) {
        /* is raw container tarball? */
        data = unzip_file(zip, FW_CONFIG, &sz);
        if (data == 0)
            die("invalid input file: %s\n", fn);
        if (parse_config(data, (size_t)sz, &conf, ver))
            die("parse config: %s failed.", FW_CONFIG);
        data = unzip_file(zip, PREOS_CONFIG, &sz);
        if (data == 0)
            die("invalid input file: %s\n", fn);
        if (parse_config(data, (size_t)sz, &conf, ver))
            die("parse config: %s failed.", PREOS_CONFIG);
    } else {
        if (parse_config(data, (size_t)sz, &conf, ver))
            die("parse config failed.");
    }

    /* basic check */
    if (strlen(conf.fwr_dnx) <= 0 || strlen(conf.ifwi) <=0 ||
            strlen(conf.boot) <= 0 || strlen(conf.preos) <= 0)
        die("invalid config file in payload.");

    data = unzip_file(zip, conf.fwr_dnx, &sz);
    if (data == 0) die("package missing %s.", conf.fwr_dnx);
    fb_queue_stream_flash("dnx", data, sz);

    data = unzip_file(zip, conf.ifwi, &sz);
    if (data == 0) die("package missing %s.", conf.ifwi);
    fb_queue_stream_flash("ifwi", data, sz);

    data = unzip_file(zip, conf.boot, &sz);
    if (data == 0) die("package missing %s.", conf.boot);
    fb_queue_stream_flash("boot", data, sz);

    data = unzip_file(zip, conf.preos, &sz);
    if (data == 0) die("package missing %s.", conf.preos);
    fb_queue_stream_flash("preos", data, sz);

    /* try to get platform image.
     * first, try gziped.
     * second, try bzip2.
     * at last, try raw image
     * otherwise, die
     */
    data = unzip_file(zip, PLATFORM_IMG ".gz", &sz);
    if (data == 0) {
        data = unzip_file(zip, PLATFORM_IMG ".bz2", &sz);
        if (data == 0) {
            data = unzip_file(zip, PLATFORM_IMG, &sz);
            if (data == 0) die("package missing platform image.");
        }
    }
    fb_queue_stream_flash("platform", data, sz);

    /* data and csa are optional */
    data = unzip_file(zip, DATA_IMG ".gz", &sz);
    if (data == 0) {
        data = unzip_file(zip, DATA_IMG ".bz2", &sz);
        if (data == 0) {
            data = unzip_file(zip, DATA_IMG, &sz);
        }
    }
    if (data != 0)
        fb_queue_stream_flash("data", data, sz);

    data = unzip_file(zip, CSA_IMG ".gz", &sz);
    if (data == 0) {
        data = unzip_file(zip, CSA_IMG ".bz2", &sz);
        if (data == 0) {
            data = unzip_file(zip, CSA_IMG, &sz);
        }
    }
    if (data != 0)
        fb_queue_stream_flash("csa", data, sz);
}

void do_send_signature(char *fn)
{
    void *data;
    unsigned sz;
    char *xtn;

    xtn = strrchr(fn, '.');
    if (!xtn) return;
    if (strcmp(xtn, ".img")) return;

    strcpy(xtn,".sig");
    data = load_file(fn, &sz);
    strcpy(xtn,".img");
    if (data == 0) return;
    fb_queue_download("signature", data, sz);
    fb_queue_command("signature", "installing signature");
}

#define skip(n) do { argc -= (n); argv += (n); } while (0)
#define require(n) do { if (argc < (n)) {usage(); exit(1);}} while (0)

int fd_pull;
char fn_pull[PATH_MAX] = "";
int do_oem_command(int argc, char **argv)
{
    int i;
    void *data;
    unsigned sz;
    char command[256];

    if (argc > 1) {
        if (0 == strcmp(argv[1], "push")) {
            if (argc > 2 && strcmp(argv[2], "-h") &&
					strcmp(argv[2], "--help")) {
                data = load_file(argv[2], &sz);
                if (data == 0) die("could not load '%s': %s", argv[2],
						strerror(errno));
                fb_queue_download(argv[2], data, sz);
            }
        } else if (0 == strcmp(argv[1], "pull")) {
            if (argc > 3)
                strncpy(fn_pull, argv[3], sizeof(fn_pull));
            else if (argc > 2 && strcmp(argv[2], "-h") &&
					strcmp(argv[2], "--help"))
                strncpy(fn_pull, basename(argv[2]), sizeof(fn_pull));
            if (strlen(fn_pull) > 0) {
                fd_pull = open(fn_pull, O_CREAT | O_EXCL | O_TRUNC | O_WRONLY,
                        S_IRUSR | S_IWUSR);
                if (fd_pull < 0) {
                    fprintf(stderr, "create local file %s: %s\n",
							fn_pull, strerror(errno));
                    return -1;
                }
            }
        }
    }

    command[0] = 0;
    while(1) {
        strcat(command,*argv);
        skip(1);
        if(argc == 0) break;
        strcat(command," ");
    }

    fb_queue_command(command,"");
    return 0;
}

int main(int argc, char **argv)
{
    int wants_reboot = 0;
    int wants_reboot_bootloader = 0;
    void *data;
    unsigned sz;
    int status;

    skip(1);
    if (argc == 0) {
        usage();
        return 1;
    }

    if (!strcmp(*argv, "-h") || !strcmp(*argv, "--help")) {
        usage();
        return 0;
    } else if (!strcmp(*argv, "-v") || !strcmp(*argv, "--version")) {
	    fprintf(stdout, "fastboot %s\n", VERSION);
	    return 0;
    }

    while (argc > 0) {
        if(!access(*argv, R_OK)) {
            /* all-in-one file */
            do_flashall(*argv);
            wants_reboot = 1;
            skip(1);
        } else if(!strcmp(*argv, "-s") || !strcmp(*argv, "--serial")) {
            require(2);
            serial = argv[1];
            skip(2);
        } else if(!strcmp(*argv, "-i") || !strcmp(*argv, "--id")) {
            char *endptr = NULL;
            unsigned long val;
            require(2);
            val = strtoul(argv[1], &endptr, 0);
            if (!endptr || *endptr != '\0' || (val & ~0xffff))
                die("invalid vendor id '%s'", argv[1]);
            vendor_id = (unsigned short)val;
            skip(2);
        } else if(!strcmp(*argv, "getvar")) {
            /* when argc == 1, just list all available variables */
            if (argc == 1) {
                fb_queue_display("","");
                skip(1);
            } else {
                require(2);
                fb_queue_display(argv[1], argv[1]);
                skip(2);
            }
        } else if(!strcmp(*argv, "erase")) {
            if (argc == 1) {
                fb_queue_erase("");
                skip(1);
            } else {
                require(2);
                fb_queue_erase(argv[1]);
                skip(2);
            }
        } else if(!strcmp(*argv, "reboot")) {
            wants_reboot = 1;
            skip(1);
        } else if(!strcmp(*argv, "reboot-bootloader")) {
            wants_reboot_bootloader = 1;
            skip(1);
        } else if(!strcmp(*argv, "flash")) {
            if (argc == 1) {
                fb_queue_stream_flash("", NULL, 0);
                skip(1);
            } else {
                char *pname = argv[1];
                char *fname = 0;
                require(3);
                fname = argv[2];
                skip(3);
                data = load_file(fname, &sz);
                if (data == 0) die("cannot load '%s': %s\n", fname, strerror(errno));
                fb_queue_stream_flash(pname, data, sz);
            }
        } else if(!strcmp(*argv, "flashall")) {
            require(2);
            do_flashall(argv[1]);
            skip(2);
            wants_reboot = 1;
        } else if(!strcmp(*argv, "oem")) {
            argc = do_oem_command(argc, argv);
            if (argc)
                return argc;
        } else if (!strcmp(*argv, "devices")) {
            list_devices();
            return 0;
        } else {
            usage();
            return 1;
        }
    }

    if (wants_reboot) {
        fb_queue_command("reboot", "rebooting");
    } else if (wants_reboot_bootloader) {
        fb_queue_command("reboot-bootloader", "rebooting into bootloader");
    }

    usb = open_device();

    status = fb_execute_queue(usb);
    return (status) ? 1 : 0;
}
