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

#include <dirent.h>
#include <getopt.h>

#include <sys/types.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 <verify.h>
#include <x509.h>

enum {
	OPTION_HELP = 'h',
	OPTION_IMAGE = 'i',
	OPTION_CERTS = 'c',
};

struct args {
	const char *image, *digest, *certs;
};

static struct opt_desc opt_descs[] = {
	{ "-i", "--image=PATH", "the image in ROTS format to verify" },
	{ "-c", "--certs=PATH", "the directory containing the certificates" },
	{ "-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 },
		{ "certs", required_argument, 0, OPTION_CERTS },
		{ NULL, 0, 0, 0 },
	};
	int ret;

	while ((ret = getopt_long(argc, (char * const *)argv, "hi:c:", options,
		NULL)) >= 0) {
		switch (ret) {
		case OPTION_HELP: return -1;
		case OPTION_IMAGE: args->image = optarg; break;
		case OPTION_CERTS: args->certs = optarg; break;
		default: return -1;
		}
	}

	if (!args->image || !args->certs)
		return -1;

	return 0;
}

static X509 *X509_find_cert_by_common_name(const char *path, const char *name)
{
	DIR *dir;
	struct dirent *dirent;
	X509 *cert = NULL;
	char *fpath, *cn;

	if (!(dir = opendir(path)))
		return NULL;

	while ((dirent = readdir(dir))) {
		if (asprintf(&fpath, "%s/%s", path, dirent->d_name) < 0)
			continue;

		if (!(cert = X509_open_cert(fpath)))
			continue;

		if (!(cn = X509_get_common_name(cert))) {
			free(cert); cert = NULL;
			continue;
		}

		if (strcmp(cn, name) == 0)
			break;
	}

	closedir(dir);

	return cert;
}

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

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

	if (!(fp = fopen(image, "rb")))
		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_DigestVerifyInit(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_DigestVerifyUpdate(ctx, data, nbytes)))
			goto err_destroy_ctx;

		size -= nbytes;
	}

	if (!(EVP_DigestVerifyUpdate(ctx, &sig_hdr->timestamp,
		sizeof sig_hdr->timestamp)))
		goto err_destroy_ctx;

	if (EVP_DigestVerifyFinal(ctx, sig, sig_len) == 1) {
		ret = 0;
	} else {
		ret = -1;
	}

	EVP_MD_CTX_destroy(ctx);
	fclose(fp);
	return ret;

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

static int verify_all(size_t *count, size_t *total, const char *ca_path, const char *image)
{
	struct rots_hdr hdr;
	struct rots_sig_hdr sig_hdr;
	unsigned char *sig;
	FILE *fp;
	EVP_PKEY *key;
	X509 *cert;
	int ret;

	*count = *total = 0;

	if (!(fp = fopen(image, "rb")))
		return -1;

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

	if (fseek(fp, hdr.size, SEEK_CUR) < 0)
		goto err_close_image;

	while (rots_read_sig_hdr(fp, &sig_hdr) == 0) {
		if (!(cert = X509_find_cert_by_common_name(ca_path, sig_hdr.name))) {
			rots_cleanup_sig_hdr(&sig_hdr);
			fseek(fp, sig_hdr.size, SEEK_CUR);
			++*total;
			continue;
		}

		if (X509_check_issued(cert, cert) == X509_V_OK) {
			printf("self-signed certificate\n");
		} else {
			printf("root-signed certificate\n");
		}

		if (!(key = X509_get_pubkey(cert))) {
			free(cert);
			rots_cleanup_sig_hdr(&sig_hdr);
			fseek(fp, sig_hdr.size, SEEK_CUR);
			++*total;
			continue;
		}

		if (!(sig = malloc(sig_hdr.size * sizeof *sig))) {
			free(cert);
			rots_cleanup_sig_hdr(&sig_hdr);
			goto err_close_image;
		}

		if (fread(sig, sizeof *sig, sig_hdr.size, fp) < sig_hdr.size) {
			free(sig);
			goto err_close_image;
		}

		ret = verify(image, sig_hdr.digest, key, &sig_hdr, sig, sig_hdr.size);
		free(sig);

		if (ret == 0)
			++*count;

		++*total;

		rots_cleanup_sig_hdr(&sig_hdr);
	}

	fclose(fp);
	return 0;

err_close_image:
	fclose(fp);
	return -1;
}

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

int do_verify(int argc, char *argv[])
{
	struct args args;
	size_t count, total;

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

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

	if (verify_all(&count, &total, args.certs, args.image) < 0) {
		fprintf(stderr, "error: unable to verify the signature(s).\n");
		return -1;
	}

	printf("%zu/%zu signatures are correct.\n", count, total);

	/* TODO: implement a configurable threshold. */
	if (count < total)
		return -1;

	return 0;
}