#include <errno.h>
#include <fcntl.h>
#include <openssl/evp.h>
#include <openssl/pem.h>
#include <openssl/rsa.h>
#include <sodium.h>
#include <string.h>
#include <sys/stat.h>
#include <unistd.h>

#include "apk_crypto.h"

/* includes '\0', undef in the end */
#define APK_CUTE_PUBLIC_SIZE 69

union public_key_state {
	unsigned char cute[crypto_sign_ed25519_PUBLICKEYBYTES];
	EVP_PKEY *pkey;
};

struct public_key_ops {
	int (*load)(struct apk_public_key *pub, int fd);
	int (*verify)(struct apk_public_key *pub, struct apk_digest_ctx *dctx, void *sig, size_t len);
	void (*free)(union public_key_state *state);
};

static int public_key_load_v2(struct apk_public_key *pub, int fd);
static int verify_rsa_pkcs1v15(struct apk_public_key *pub, struct apk_digest_ctx *dctx, void *sig, size_t len);
static void public_key_free_v2(union public_key_state *state);

static int public_key_load_cute(struct apk_public_key *pub, int fd);
static int verify_cute(struct apk_public_key *pub, struct apk_digest_ctx *dctx, void *sig, size_t len);
static void public_key_free_cute(union public_key_state *state);

static const struct public_key_ops ops[] = {
    [APK_KEY_VERSION_V2]   = {.load = public_key_load_v2, .verify = verify_rsa_pkcs1v15, .free = public_key_free_v2},
    [APK_KEY_VERSION_CUTE] = {.load = public_key_load_cute, .verify = verify_cute, .free = public_key_free_cute},
};

/*
 * The format is as follows:
 *     uint8[ 2] header = {'q', 't'}
 *     uint8[16] key_id
 *     uint8[32] public_key
 */
static int public_key_load_cute(struct apk_public_key *pub, int fd)
{
	union public_key_state *state = (union public_key_state *) pub->impl;

	char b64[APK_CUTE_PUBLIC_SIZE];
	uint8_t raw[50];
	size_t len;

	if (read(fd, b64, APK_CUTE_PUBLIC_SIZE) != APK_CUTE_PUBLIC_SIZE) {
		return -errno;
	}

	if (sodium_base642bin(raw, 50, b64, APK_CUTE_PUBLIC_SIZE, NULL, &len, NULL, sodium_base64_VARIANT_ORIGINAL)
	    != 0) {
		return -APKE_CRYPTO_KEY_FORMAT;
	}

	if (len != 50) {
		return -APKE_CRYPTO_KEY_FORMAT;
	}

	if (raw[0] != 'q' || raw[1] != 't') {
		return -APKE_CRYPTO_KEY_FORMAT;
	}

	memcpy(pub->id, raw + 2, sizeof pub->id);
	memcpy(state->cute, raw + 18, crypto_sign_ed25519_PUBLICKEYBYTES);

	return 0;
}

/*
 * The signature format is as follows:
 *     uint8[ 2] header = {'q', 't'}
 *     uint8[16] key_id
 *     uint8[64] signature
 */
static int verify_cute(struct apk_public_key *pub, struct apk_digest_ctx *dctx, void *sig, size_t len)
{
	union public_key_state *state = (union public_key_state *) pub->impl;

	unsigned char *sigdata = (unsigned char *) sig;
	struct apk_digest digest;

	if (len != 82) {
		return -APKE_SIGNATURE_INVALID;
	}

	if (sigdata[0] != 'q' || sigdata[1] != 't') {
		return -APKE_SIGNATURE_INVALID;
	}

	if (sodium_memcmp(pub->id, sigdata + 2, sizeof pub->id) != 0) {
		return -APKE_SIGNATURE_INVALID;
	}

	if (dctx->alg != APK_DIGEST_SHA256) {
		return -APKE_SIGNATURE_INVALID;
	}

	if (apk_digest_ctx_final(dctx, &digest) != 0) {
		return -APKE_SIGNATURE_INVALID;
	}

	if (crypto_sign_ed25519_verify_detached(sigdata + 18, digest.data, digest.len, state->cute) != 0) {
		return -APKE_SIGNATURE_INVALID;
	}

	return 0;
}

static void public_key_free_cute(union public_key_state *state)
{
	return;
}

static int public_key_load_v2(struct apk_public_key *pub, int fd)
{
	union public_key_state *state = (union public_key_state *) pub->impl;

	unsigned char *raw = NULL;
	struct apk_digest digest;
	BIO *bio;
	int len;
	int r = -APKE_CRYPTO_KEY_FORMAT;

	bio = BIO_new_fd(fd, BIO_NOCLOSE);
	if (bio == NULL) {
		r = -ENOMEM;
		goto err_bio_new;
	}

	state->pkey = PEM_read_bio_PUBKEY(bio, NULL, NULL, NULL);
	if (state->pkey == NULL) {
		r = -APKE_CRYPTO_KEY_FORMAT;
		goto err_load_and_check_type;
	}

	if (state->pkey == NULL) {
		r = -APKE_CRYPTO_KEY_FORMAT;
		goto err_load_and_check_type;
	}

	if (EVP_PKEY_id(state->pkey) != EVP_PKEY_RSA) {
		r = -APKE_CRYPTO_KEY_FORMAT;
		goto err_load_and_check_type;
	}

	len = i2d_PublicKey(state->pkey, &raw);
	if (len < 0) {
		r = -APKE_CRYPTO_ERROR;
		goto err_load_and_check_type;
	}

	if (apk_digest_calc(&digest, APK_DIGEST_SHA512, raw, len) != 0) {
		r = -APKE_CRYPTO_ERROR;
		goto err_calculate_id;
	}

	memcpy(pub->id, digest.data, sizeof pub->id);

	r = 0;

err_calculate_id:
	OPENSSL_free(raw);
err_load_and_check_type:
	BIO_free(bio);
err_bio_new:
	return r;
}

/*
 * RSA-PKCS1 v1.5 can work with different hashes.
 * Since APK v2 hasn't switched from SHA-1 yet, we still have to work with it for now.
 * However, MD5 is out of the question.
 */
static int verify_rsa_pkcs1v15(struct apk_public_key *pub, struct apk_digest_ctx *dctx, void *sig, size_t len)
{
	union public_key_state *state = (union public_key_state *) pub->impl;

	struct apk_digest digest;
	EVP_PKEY_CTX *pctx;
	const EVP_MD *md;
	int r = -APKE_SIGNATURE_INVALID;

	if (apk_digest_ctx_final(dctx, &digest) != 0) {
		r = -APKE_SIGNATURE_INVALID;
		goto err_digest_final_and_pctx_new;
	}

	switch (dctx->alg) {
	case APK_DIGEST_SHA1:
		md = EVP_sha1();
		break;
	case APK_DIGEST_SHA256:
		md = EVP_sha256();
		break;
	case APK_DIGEST_SHA512:
		md = EVP_sha512();
		break;
	case APK_DIGEST_NONE:
	case APK_DIGEST_MD5:
	default:
		r = -APKE_SIGNATURE_INVALID;
		goto err_digest_final_and_pctx_new;
	}

	pctx = EVP_PKEY_CTX_new(state->pkey, NULL);
	if (pctx == NULL) {
		r = -APKE_SIGNATURE_INVALID;
		goto err_pctx_setup_and_verify;
	}

	if (EVP_PKEY_verify_init(pctx) != 1) {
		r = -APKE_SIGNATURE_INVALID;
		goto err_pctx_setup_and_verify;
	}

	if (EVP_PKEY_CTX_set_rsa_padding(pctx, RSA_PKCS1_PADDING) != 1) {
		r = -APKE_SIGNATURE_INVALID;
		goto err_pctx_setup_and_verify;
	}

	if (EVP_PKEY_CTX_set_signature_md(pctx, md) != 1) {
		r = -APKE_SIGNATURE_INVALID;
		goto err_pctx_setup_and_verify;
	}

	if (EVP_PKEY_verify(pctx, sig, len, digest.data, digest.len) != 1) {
		r = -APKE_SIGNATURE_INVALID;
		goto err_pctx_setup_and_verify;
	}

	r = 0;

err_pctx_setup_and_verify:
	EVP_PKEY_CTX_free(pctx);
err_digest_final_and_pctx_new:
	return r;
}

static void public_key_free_v2(union public_key_state *state)
{
	EVP_PKEY_free(state->pkey);
}

int apk_public_key_load(struct apk_public_key *pub, int dirfd, const char *fn)
{
	struct stat st;
	int fd = -1;
	int r  = -APKE_CRYPTO_KEY_FORMAT;

	pub->impl = malloc(sizeof(union public_key_state));
	if (pub->impl == NULL) {
		r = -ENOMEM;
		goto err_alloc;
	}

	fd = openat(dirfd, fn, O_RDONLY | O_CLOEXEC);
	if (fd == -1) {
		r = -errno;
		goto err_openat;
	}

	if (fstat(fd, &st) != 0) {
		r = -errno;
		goto err_fstat_and_load;
	}

	// guess
	if (st.st_size == APK_CUTE_PUBLIC_SIZE) {
		pub->version = APK_KEY_VERSION_CUTE;
	} else {
		pub->version = APK_KEY_VERSION_V2;
	}

	r = ops[pub->version].load(pub, fd);
	if (r != 0) {
		goto err_fstat_and_load;
	}

	close(fd);

	return 0;

err_fstat_and_load:
	close(fd);
err_openat:
	free(pub->impl);
err_alloc:
	return r;
}

void apk_public_key_free(struct apk_public_key *pub)
{
	if (pub == NULL) {
		return;
	}

	if (pub->impl == NULL || pub->version >= APK_KEY_VERSION_MAX) {
		return;
	}

	ops[pub->version].free(pub->impl);

	free(pub->impl);
	pub->impl = NULL;
}

int apk_verify_digest_start(struct apk_digest_ctx *dctx, uint16_t signature_type)
{
	uint8_t digest;

	switch (signature_type) {
	case APK_SIGNATURE_CUTE:
	case APK_SIGNATURE_RSA256:
		digest = APK_DIGEST_SHA256;
		break;
	case APK_SIGNATURE_RSA512:
		digest = APK_DIGEST_SHA512;
		break;
	case APK_SIGNATURE_RSA:
		digest = APK_DIGEST_SHA1;
		break;
	default:
		return -APKE_CRYPTO_NOT_SUPPORTED;
	}

	if (apk_digest_ctx_reset(dctx, digest) != 0) {
		return -APKE_CRYPTO_ERROR;
	}

	return 0;
}

int apk_verify(struct apk_public_key *pub, struct apk_digest_ctx *dctx, void *sig, size_t len)
{
	if (pub->version >= APK_KEY_VERSION_MAX) {
		return -APKE_CRYPTO_NOT_SUPPORTED;
	}

	return ops[pub->version].verify(pub, dctx, sig, len);
}

#undef APK_CUTE_PUBLIC_SIZE