#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

#include <getopt.h>

#include <openssl/err.h>
#include <openssl/evp.h>
#include <openssl/pem.h>
#include <openssl/ssl.h>
#include <openssl/x509.h>
#include <openssl/x509v3.h>

#include <image.h>
#include <macros.h>
#include <option.h>
#include <sign.h>
#include <x509.h>

enum {
	OPTION_HELP = 'h',
	OPTION_IMAGE = 'i',
	OPTION_DIGEST = 'd',
	OPTION_KEY = 'k',
	OPTION_CERT = 'c',
};

struct args {
	const char *image, *digest, *key, *cert;
};

static struct opt_desc opt_descs[] = {
	{ "-i", "--image=PATH","the image in ROTS format to sign" },
	{ "-d", "--digest={ripemd160,sha224,sha256,sha384,sha512,whirlpool}\n",
		"\t\tthe digest algorithm to use for signing" },
	{ "-k", "--key=PATH", "the private key" },
	{ "-c", "--cert=PATH", "the certificate containing the public key" },
	{ "-h", "--help", "display this help and exit" },
	{ NULL, NULL, NULL },
};

static int parse_args(struct args *args, int argc, char *argv[])
{
	struct option options[] = {
		{ "help", no_argument, NULL, OPTION_HELP },
		{ "image", required_argument, 0, OPTION_IMAGE },
		{ "digest", required_argument, 0, OPTION_DIGEST },
		{ "key", required_argument, 0, OPTION_KEY },
		{ "cert", required_argument, 0, OPTION_CERT },
		{ NULL, 0, 0, 0 },
	};
	int ret;

	while ((ret = getopt_long(argc, (char * const *)argv, "hi:d:k:c:", options,
		NULL)) >= 0) {
		switch (ret) {
		case OPTION_HELP: return -1;
		case OPTION_IMAGE: args->image = optarg; break;
		case OPTION_DIGEST: args->digest = optarg; break;
		case OPTION_KEY: args->key = optarg; break;
		case OPTION_CERT: args->cert = optarg; break;
		default: return -1;
		}
	}

	if (!args->image || !args->digest || !args->key || !args->cert)
		return -1;

	return 0;
}

static EVP_PKEY *open_priv_key(const char *path)
{
	EVP_PKEY *key;
	FILE *fp;

	if (!(fp = fopen(path, "r")))
		return NULL;

	PEM_read_PrivateKey(fp, &key, NULL, NULL);
	fclose(fp);

	return key;
}

static int sign(const char *image, const char *name, const char *digest_name,
	EVP_PKEY *key)
{
	char data[512];
	struct rots_hdr hdr;
	struct rots_sig_hdr sig_hdr;
	const EVP_MD *digest;
	FILE *fp;
	unsigned char *sig;
	EVP_MD_CTX *ctx = NULL;
	size_t nbytes, size, sig_len = 0;

	if (!image || !digest_name || !key)
		return -1;

	if (!(fp = fopen(image, "r+b")))
		return -1;

	if (rots_read_hdr(fp, &hdr) < 0)
		goto err_close_image;

	size = hdr.size + ftell(fp);
	fseek(fp, 0, SEEK_SET);

	if (!(ctx = EVP_MD_CTX_create()))
		goto err_close_image;

	if (!(digest = EVP_get_digestbyname(digest_name)))
		goto err_destroy_ctx;

	if (!(EVP_DigestSignInit(ctx, NULL, digest, NULL, key)))
		goto err_destroy_ctx;

	while (size) {
		nbytes = fread(data, sizeof *data, min(size, sizeof data), fp);

		if (nbytes == 0)
			goto err_destroy_ctx;

		if (!(EVP_DigestSignUpdate(ctx, data, nbytes)))
			goto err_destroy_ctx;

		size -= nbytes;
	}

	sig_hdr.timestamp = (uint64_t)time(NULL);

	if (!(EVP_DigestSignUpdate(ctx, &sig_hdr.timestamp,
		sizeof sig_hdr.timestamp)))
		goto err_destroy_ctx;

	if (!(EVP_DigestSignFinal(ctx, NULL, &sig_len)))
		goto err_destroy_ctx;

	if (!(sig = malloc(sizeof(char) * sig_len)))
		goto err_destroy_ctx;

	if (!(EVP_DigestSignFinal(ctx, sig, &sig_len)))
		goto err_free_sig;

	if (fseek(fp, 0, SEEK_END) < 0)
		goto err_free_sig;

	sig_hdr.name = name;
	sig_hdr.digest = digest_name;
	sig_hdr.size = sig_len;

	if (rots_write_sig_hdr(fp, &sig_hdr) < 0)
		goto err_free_sig;

	if (fwrite(sig, sizeof *sig, sig_len, fp) < sig_len)
		goto err_free_sig;

	free(sig);
	EVP_MD_CTX_destroy(ctx);
	fclose(fp);

	return 0;

err_free_sig:
	free(sig);
err_destroy_ctx:
	EVP_MD_CTX_destroy(ctx);
err_close_image:
	fclose(fp);
	return -1;
}

void show_sign_usage(const char *prog_name, const char *cmd)
{
	fprintf(stderr, "usage: %s %s [option]...\n\n"
		"sign an image\n\n", prog_name, cmd);
	format_options(opt_descs);
}

int do_sign(int argc, char *argv[])
{
	struct args args;
	EVP_PKEY *key;
	X509 *cert;
	char *cn;

	if (parse_args(&args, argc, argv) < 0) {
		show_sign_usage(argv[0], argv[1]);
		return -1;
	}

	SSL_load_error_strings();
	SSL_library_init();
	OpenSSL_add_all_algorithms();

	if (!(key = open_priv_key(args.key))) {
		fprintf(stderr, "error: unable to read the private key.\n");
		ERR_print_errors_fp(stderr);
		return -1;
	}

	if (!(cert = X509_open_cert(args.cert))) {
		fprintf(stderr, "error: unable to read the X509 certificate.\n");
		ERR_print_errors_fp(stderr);
		OPENSSL_free(key);
		return -1;
	}

	if (!(cn = X509_get_common_name(cert))) {
		fprintf(stderr, "error: unable to get the common name.\n");
		ERR_print_errors_fp(stderr);
		OPENSSL_free(cert);
		OPENSSL_free(key);
		return -1;
	}
	
	if (sign(args.image, cn, args.digest, key) < 0) {
		fprintf(stderr, "error: unable to sign the payload.\n");
		ERR_print_errors_fp(stderr);
		OPENSSL_free(cn);
		OPENSSL_free(cert);
		OPENSSL_free(key);
		return -1;
	}

	OPENSSL_free(cn);
	OPENSSL_free(cert);
	OPENSSL_free(key);

	return 0;
}