#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_SECRET_SIZE 113

union secret_key_state {
	unsigned char cute[crypto_sign_ed25519_SECRETKEYBYTES];
	EVP_PKEY *pkey;
};

struct secret_key_ops {
	int (*load)(struct apk_secret_key *pub, int fd);
	int (*sign)(struct apk_secret_key *pub, struct apk_digest_ctx *dctx, void *sig, size_t *len);
	void (*free)(union secret_key_state *state);
};

static int secret_key_load_v2(struct apk_secret_key *pub, int fd);
static int sign_rsa_pkcs1v15(struct apk_secret_key *pub, struct apk_digest_ctx *dctx, void *sig, size_t *len);
static void secret_key_free_v2(union secret_key_state *state);

static int secret_key_load_cute(struct apk_secret_key *pub, int fd);
static int sign_cute(struct apk_secret_key *pub, struct apk_digest_ctx *dctx, void *sig, size_t *len);
static void secret_key_free_cute(union secret_key_state *state);

static struct secret_key_ops ops[] = {
    [APK_KEY_VERSION_CUTE] = {.load = secret_key_load_cute, .sign = sign_cute, .free = secret_key_free_cute},
    [APK_KEY_VERSION_V2]   = {.load = secret_key_load_v2, .sign = sign_rsa_pkcs1v15, .free = secret_key_free_v2},
};

/*
 * The format is as follows:
 *     uint8[ 2] header = {'q', 't'}
 *     uint8[16] key_id
 *     uint8[64] secret_key
 *
 */
static int secret_key_load_cute(struct apk_secret_key *sec, int fd)
{
	union secret_key_state *state = (union secret_key_state *) sec->impl;

	char b64[APK_CUTE_SECRET_SIZE];
	uint8_t raw[82];
	size_t len;
	int r = -APKE_CRYPTO_KEY_FORMAT;

	if (read(fd, b64, APK_CUTE_SECRET_SIZE) != APK_CUTE_SECRET_SIZE) {
		r = -errno;
		goto err;
	}

	if (sodium_base642bin(raw, 82, b64, APK_CUTE_SECRET_SIZE, NULL, &len, NULL, sodium_base64_VARIANT_ORIGINAL)
	    != 0) {
		r = -APKE_CRYPTO_KEY_FORMAT;
		goto err;
	}

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

	sec->version = APK_KEY_VERSION_V2;
	memcpy(sec->id, raw + 2, 16);
	memcpy(state->cute, raw + 18, crypto_sign_ed25519_SECRETKEYBYTES);

	r = 0;

err:
	sodium_memzero(b64, sizeof b64);
	sodium_memzero(raw, sizeof raw);
	return r;
}

static void secret_key_free_cute(union secret_key_state *state)
{
	sodium_memzero(state->cute, crypto_sign_ed25519_SECRETKEYBYTES);
}

static int sign_cute(struct apk_secret_key *sec, struct apk_digest_ctx *dctx, void *sig, size_t *len)
{
	union secret_key_state *state = (union secret_key_state *) sec->impl;

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

	if (*len < 82) {
		return -APKE_SIGNATURE_FAIL;
	}

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

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

	sigdata[0] = 'q';
	sigdata[1] = 't';
	memcpy(sigdata + 2, sec->id, 16);
	crypto_sign_ed25519_detached(sigdata + 18, NULL, digest.data, digest.len, state->cute);
	*len = 82;

	return 0;
}

static int secret_key_load_v2(struct apk_secret_key *sec, int fd)
{
	union secret_key_state *state = (union secret_key_state *) sec->impl;

	struct apk_digest digest;
	unsigned char *raw = NULL;
	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_PrivateKey(bio, NULL, NULL, NULL);
	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_digest;
	}

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

	r = 0;

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

static int sign_rsa_pkcs1v15(struct apk_secret_key *sec, struct apk_digest_ctx *dctx, void *sig, size_t *len)
{
	union secret_key_state *state = (union secret_key_state *) sec->impl;

	struct apk_digest digest;
	EVP_PKEY_CTX *pctx;
	const EVP_MD *md;
	int r = -APKE_SIGNATURE_FAIL;
	size_t needed;

	if (apk_digest_ctx_final(dctx, &digest) != 0) {
		r = -APKE_SIGNATURE_FAIL;
		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_FAIL;
		goto err_digest_final_and_pctx_new;
	}

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

	if (EVP_PKEY_sign_init(pctx) != 1) {
		r = -APKE_SIGNATURE_FAIL;
		goto err_pctx_setup_and_sign;
	}

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

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

	/* Determine minimum required buffer size */
	if (EVP_PKEY_sign(pctx, NULL, &needed, digest.data, digest.len) != 1) {
		r = -APKE_SIGNATURE_FAIL;
		goto err_pctx_setup_and_sign;
	}

	if (needed < *len) {
		r = -APKE_SIGNATURE_FAIL;
		goto err_pctx_setup_and_sign;
	}

	if (EVP_PKEY_sign(pctx, sig, len, digest.data, digest.len) != 1) {
		r = -APKE_SIGNATURE_FAIL;
		goto err_pctx_setup_and_sign;
	}

	r = 0;

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

static void secret_key_free_v2(union secret_key_state *state)
{
	EVP_PKEY_free(state->pkey);
}

int apk_secret_key_load(struct apk_secret_key *sec, int dirfd, const char *fn)
{
	struct stat st;
	int fd = -1;
	int r  = -APKE_CRYPTO_KEY_FORMAT;

	sec->impl = malloc(sizeof(union secret_key_state));
	if (sec->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;
	}

	if (st.st_size == APK_CUTE_SECRET_SIZE) {
		sec->version = APK_KEY_VERSION_CUTE;
	} else {
		sec->version = APK_KEY_VERSION_V2;
	}

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

	close(fd);

	return 0;

err_load:
err_fstat:
	close(fd);
err_openat:
	free(sec->impl);
err_alloc:
	return r;
}

void apk_secret_key_free(struct apk_secret_key *sec)
{
	if (sec == NULL) {
		return;
	}

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

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

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

int apk_sign_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_sign(struct apk_secret_key *sec, struct apk_digest_ctx *dctx, void *sig, size_t *len)
{
	if (sec->version >= APK_KEY_VERSION_MAX) {
		return -APKE_CRYPTO_NOT_SUPPORTED;
	}

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

#undef APK_CUTE_SECRET_SIZE