/*
    volez - generate update.zip for Android 2.0 / 2.0.1 with arbitrary data
    Copyright (C) 2009  Christopher Lais

    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License along
    with this program; if not, write to the Free Software Foundation, Inc.,
    51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
*/

#include <iostream>
#include <climits>
#include <fstream>

#define VERSION	"2009-12-11"

extern "C" {
#include <sys/types.h>
#include <unistd.h>
#include <fcntl.h>
#if !WINDOWS
#include <sys/mman.h>
#include <fnmatch.h>
#define O_BINARY	0 /* Only Windows is brain-damaged. */
#else
#include "win32/mingw_mmap.h"
#include "win32/fnmatch.h"
#endif
};

#include <map>
#include <set>

#include "zip.h"

#define PAGE_ALIGN(x) (((x) + PAGE_SIZE - 1) & ~(PAGE_SIZE - 1))

static unsigned short le16(const unsigned char *mem)
{
	return mem[0] | (mem[1]<<8);
}

static unsigned int le32(const unsigned char *mem)
{
	return mem[0] | (mem[1]<<8) | (mem[2]<<16) | (mem[3]<<24);
}

static void write_le16(unsigned char *mem, unsigned short val)
{
	mem[0] = (val >> 0) & 0xff;
	mem[1] = (val >> 8) & 0xff;
}

static void write_le32(unsigned char *mem, unsigned int val)
{
	mem[0] = (val >> 0) & 0xff;
	mem[1] = (val >> 8) & 0xff;
	mem[2] = (val >> 16) & 0xff;
	mem[3] = (val >> 24) & 0xff;
}

struct zip_file_t {
	unsigned char *mem;
	unsigned char *eocd;
	size_t len;
};

struct zip_entry_t {
	off_t lf_ofs;
	size_t lf_len;
	off_t cd_ofs;
	size_t cd_len;
};

#define RSANUMBYTES         256

#define FOOTER_MAGIC	0xffff
#define FOOTER_SIG_OFS      0
#define FOOTER_MAGIC_OFS    2
#define FOOTER_COMMENT_OFS  4
#define FOOTER_LEN          6
struct signapk_sig_t {
	unsigned char *signature;
	unsigned char *footer;
};

typedef std::map<std::string, zip_entry_t> zip_dir_t;
typedef std::map<std::string, bool> zip_match_t;

unsigned char *zip_eocd(zip_file_t *zip)
{
	for (unsigned char *eocd = zip->mem + zip->len - 4; eocd >= zip->mem; --eocd) {
		if (*eocd == (EOCD_MAGIC & 0xff) && le32(eocd) == EOCD_MAGIC)
			return eocd;
	}
	return NULL;
}

int zip_open(zip_file_t *zip, int fd)
{
	zip->len = lseek(fd, 0, SEEK_END);
	if (zip->len < 0) return -1;

	zip->mem = (unsigned char*)mmap(NULL, zip->len, PROT_READ|PROT_WRITE, MAP_PRIVATE, fd, 0);
	if (zip->mem == (unsigned char*)((void*)-1)) return -1;

	zip->eocd = zip_eocd(zip);
	if (zip->eocd == NULL) return -1;

	return 0;
}

int zip_open(zip_file_t *zip, const char *filename)
{
	int fd = open(filename, O_RDONLY|O_BINARY);
	if (fd < 0) {
		std::cerr << "Unable to open " << filename << std::endl;
		return -1;
	}

	if (zip_open(zip, fd)) {
		std::cerr << "Unable to load " << filename << std::endl;
		return -1;
	}

	close(fd);

	return 0;
}

int zip_cd_load(zip_file_t *zip, zip_dir_t &dir, bool default_include = true, zip_match_t *glob_match = NULL, zip_dir_t *exclude = NULL)
{
	unsigned int cd_ofs = le32(zip->eocd + EOCD_CD_OFS);
	if (cd_ofs > (zip->eocd - zip->mem)) {
		return -1;
	}

	unsigned char *cd = zip->mem + cd_ofs;
	while (cd < zip->eocd) {
		zip_entry_t entry;

		if (le32(cd) != CDH_MAGIC) {
			std::cerr << "Bad central directory magic" << std::endl;
			return -1;
		}

		unsigned short filename_len = le16(cd + CDH_FNLEN_OFS);
		unsigned short ex_len = le16(cd + CDH_EXLEN_OFS);
		unsigned short comment_len = le16(cd + CDH_COMLEN_OFS);
		unsigned int lfh_ofs = le32(cd + CDH_LFH_OFS);

		entry.lf_ofs = lfh_ofs;
		if (entry.lf_ofs > (zip->eocd - zip->mem)) {
			std::cerr << "Local file header reads past end of zip" << std::endl;
			return -1;
		}
		unsigned char *lfh = zip->mem + entry.lf_ofs;
		if (le32(lfh) != LFH_MAGIC) {
			std::cerr << "Bad local file magic" << std::endl;
			return -1;
		}

		entry.lf_len = LFH_FILENAME_OFS + le16(lfh + LFH_FNLEN_OFS) + le16(lfh + LFH_EXLEN_OFS) + le32(lfh + LFH_CSIZE_OFS);
		if ((lfh + entry.lf_len) > zip->eocd) {
			std::cerr << "Local file data reads past end of zip" << std::endl;
			return -1;
		}

		entry.cd_ofs = cd - zip->mem;
		entry.cd_len = CDH_FILENAME_OFS + filename_len + ex_len + comment_len;
		if ((cd + entry.cd_len) > zip->eocd) {
			std::cerr << "Central directory reads past end of zip" << std::endl;
			return -1;
		}

		std::string filename = std::string((const char*)(cd + CDH_FILENAME_OFS), filename_len);

		int match_len = 0;
		bool match_include = default_include;
		if (glob_match) {
			/* Could speed this up by sorting by size. */
			for (zip_match_t::const_iterator i = glob_match->begin(); i != glob_match->end(); ++i) {
				if (i->first.size() <= match_len) continue;
				if (fnmatch(i->first.c_str(), filename.c_str(), 0)) continue;
				match_len = i->first.size();
				match_include = i->second;
			}
		}

		if (exclude) {
			if (exclude->find(filename) != exclude->end())
				match_include = false;
		}

		if (match_include) {
			dir[filename] = entry;
		}

		cd += entry.cd_len;
	}

	return 0;
}

int zip_sig_load(zip_file_t *zip, signapk_sig_t *sig)
{
	sig->footer = zip->mem + zip->len - FOOTER_LEN;

	if (le16(sig->footer + FOOTER_MAGIC_OFS) != FOOTER_MAGIC) {
		return -1;
	}

	unsigned short signature_len = le16(sig->footer + FOOTER_SIG_OFS);
	unsigned short comment_len = le16(sig->footer + FOOTER_COMMENT_OFS);

	if (comment_len < signature_len) {
		std::cerr << "SignApk signature smaller than zip comment" << std::endl;
		return -1;
	}

	if (signature_len < RSANUMBYTES) {
		std::cerr << "SignApk signature too small" << std::endl;
		return -1;
	}

	if ((zip->mem + zip->len - comment_len - EOCD_COMMENT_OFS) != zip->eocd) {
		std::cerr << "EOCD not before SignApk signature" << std::endl;
		return -1;
	}

	sig->signature = sig->footer - RSANUMBYTES;

	return 0;
}

static void show_help(char *bin)
{
	std::cout
		<< "    volez " VERSION ", Copyright (C) 2009 Christopher Lais" << std::endl
		<< "    volez comes with ABSOLUTELY NO WARRANTY." << std::endl
		<< std::endl;

	std::cout << bin << " [-I ...] [-x ...] [-o ...] signed-update.zip payload.zip" << std::endl;
	std::cout << "\t-I pattern     Include files from signed-update.zip" << std::endl;
	std::cout << "\t               matching the specified pattern." << std::endl;
	std::cout << "\t               -I can be given multiple times." << std::endl;
	std::cout << "\t-x pattern     Exclude files from signed-update.zip" << std::endl;
	std::cout << "\t               matching the specified pattern." << std::endl;
	std::cout << "\t               -x can be given multiple times." << std::endl;
	std::cout << "\t       You may specify overlapping patterns." << std::endl;
	std::cout << "\t       If you do, the longest pattern will be used." << std::endl;
	std::cout << "\t       Files in payload.zip are ALWAYS included" << std::endl;
	std::cout << "\t       from payload.zip." << std::endl;
	std::cout << "\t-o output      Write the output to the specified file." << std::endl;
}

#define COMMENT	"signed by SignApk\0Created by Volez " VERSION "; http://zenthought.org/content/project/volez"
#define COMMENT_LEN (sizeof COMMENT)

int main(int argc, char *argv[])
{
	char *donor, *payload, *output = NULL;
	zip_match_t glob_match;

	int ch;
	while ((ch = getopt(argc, argv, "I:x:o:h")) != -1) {
		switch (ch) {
		case 'I':
			glob_match[optarg] = true;
			break;
		case 'x':
			glob_match[optarg] = false;
			break;
		case 'o':
			output = optarg;
			break;
		case 'h':
			show_help(argv[0]);
			return 0;
		default:
			std::cerr << "Bad argument; use -h for help." << std::endl;
			return -1;
		}
	}

	if ((argc - optind) < 2) {
		if (argc == 1) { show_help(argv[0]); }
		std::cerr << "Missing argument; use -h for help." << std::endl;
		return -1;
	}

	std::cerr
		<< "    volez " VERSION ", Copyright (C) 2009 Christopher Lais" << std::endl
		<< "    volez comes with ABSOLUTELY NO WARRANTY." << std::endl
		<< std::endl;

	donor = argv[optind];
	payload = argv[optind + 1];

	zip_file_t donor_zip, payload_zip;
	zip_dir_t donor_dir, payload_dir;
	signapk_sig_t sig;

	int output_fd = output ? open(output, O_WRONLY|O_CREAT|O_TRUNC|O_BINARY, 0666) : STDOUT_FILENO;
	if (output && output_fd < 0) {
		std::cerr << "Unable to open " << output << " for output" << std::endl;
		return -1;
	}
	if (isatty(output_fd)) {
		show_help(argv[0]);
		std::cerr << "Refusing output to tty." << std::endl;
		return -1;
	}

	if (zip_open(&donor_zip, donor)) {
		return -1;
	}

	if (zip_open(&payload_zip, payload)) {
		return -1;
	}

	if (zip_cd_load(&payload_zip, payload_dir)) {
		std::cerr << "Unable to load central directory from " << payload << std::endl;
		return -1;
	}

	if (zip_cd_load(&donor_zip, donor_dir, false, &glob_match, &payload_dir)) {
		std::cerr << "Unable to load central directory from " << donor << std::endl;
		return -1;
	}

	if (zip_sig_load(&donor_zip, &sig)) {
		std::cerr << "Unable to read SignApk signature from " << donor << std::endl;
		return -1;
	}

	off_t offset = donor_zip.len;
	unsigned int files = 0;

	/* Write out the payload local files and update offset to final zip */
	for (zip_dir_t::iterator file = payload_dir.begin(); file != payload_dir.end(); ++file) {
		write(output_fd, payload_zip.mem + file->second.lf_ofs, file->second.lf_len);
		file->second.lf_ofs = offset;
		offset += file->second.lf_len;
	}

	off_t cd_ofs = offset;

	/* Write central directory entries for signed zip */
	for (zip_dir_t::iterator file = donor_dir.begin(); file != donor_dir.end(); ++file) {
		write(output_fd, donor_zip.mem + file->second.cd_ofs, file->second.cd_len);
		offset += file->second.cd_len;
		++files;
	}

	/* Write central directory entries for payload zip with updated offsets for final zip */
	for (zip_dir_t::iterator file = payload_dir.begin(); file != payload_dir.end(); ++file) {
		unsigned char *cd = payload_zip.mem + file->second.cd_ofs;
		write_le32(cd + CDH_LFH_OFS, file->second.lf_ofs);
		write(output_fd, cd, file->second.cd_len);
		offset += file->second.cd_len;
		++files;
	}

	off_t cd_size = offset - cd_ofs;

	/* Rebuild/Write the EOCD */
	unsigned char *eocd = payload_zip.eocd;
	write_le16(eocd + EOCD_FILES_OFS, files);
	write_le16(eocd + EOCD_CDFILES_OFS, files);
	write_le32(eocd + EOCD_SIZE_OFS, cd_size);
	write_le32(eocd + EOCD_CD_OFS, cd_ofs);
	write_le16(eocd + EOCD_COMLEN_OFS, COMMENT_LEN + RSANUMBYTES + FOOTER_LEN);
	write(output_fd, eocd, EOCD_COMMENT_OFS);
	offset += EOCD_COMMENT_OFS;

	write(output_fd, COMMENT, COMMENT_LEN);
	offset += COMMENT_LEN;

	write(output_fd, sig.signature, RSANUMBYTES);
	offset += RSANUMBYTES;

	/* Write the SignApk footer */
	offset += FOOTER_LEN;

	unsigned int footer_sig_len = le16(sig.footer + FOOTER_SIG_OFS) + offset - donor_zip.len;
	unsigned int footer_comment_len = le16(sig.footer + FOOTER_COMMENT_OFS) + offset - donor_zip.len;

	if ((footer_sig_len & ~0xffff) || (footer_comment_len & ~0xffff)) {
		std::cerr << "Output too large, make the payload smaller." << std::endl;
		return -1;
	}

	write_le16(sig.footer + FOOTER_SIG_OFS, footer_sig_len);
	write_le16(sig.footer + FOOTER_COMMENT_OFS, footer_comment_len);
	write(output_fd, sig.footer, FOOTER_LEN);

	/* All done */
	close(output_fd);

#if !WINDOWS
	std::cerr << "Created; 'cat " << donor << " " << (output?output:"<output_zip>") << " > update.zip' to create final update.zip" << std::endl;
#else
	std::cerr << "Created; 'copy /b " << donor << "+" << (output?output:"<output_zip>") << " update.zip' to create final update.zip" << std::endl;
#endif

	return 0;
}
