#include "ecc_api.h"

#include <openssl/evp.h>
#include <openssl/err.h>
#include <openssl/ec.h>
#include <openssl/pem.h>
#include <openssl/sha.h>
#include <openssl/ecdsa.h>
#include <string.h>

static const int _NIDS[] = {
	NID_secp256k1,
	NID_secp256k1
};

static EVP_PKEY_CTX *pctx[sizeof(_NIDS) / sizeof(_NIDS[0])] = { NULL };
static EVP_PKEY_CTX *kctx[sizeof(_NIDS) / sizeof(_NIDS[0])] = { NULL };
static EVP_PKEY     *params[sizeof(_NIDS) / sizeof(_NIDS[0])] = { NULL };

int ecc_api_init()
{
	int i;
	int rc = -1;
	for (i = 0; i < sizeof(_NIDS) / sizeof(_NIDS[0]); i++){
		/* Create the context for generating the parameters */
		pctx[i] = EVP_PKEY_CTX_new_id(EVP_PKEY_EC, NULL);
		if (pctx[i]){
			if (EVP_PKEY_paramgen_init(pctx[i])){
				/* Set the paramgen parameters according to the type */
				if (EVP_PKEY_CTX_set_ec_paramgen_curve_nid(pctx[i], _NIDS[i])){
					/* Generate parameters */
					if (EVP_PKEY_paramgen(pctx[i], &params[i])) {
						/* Create context for the key generation */
						kctx[i] = EVP_PKEY_CTX_new(params[i], NULL);
						if (kctx[i]){
							if (EVP_PKEY_keygen_init(kctx[i])){
								rc = 0;
								continue;
							}
						}
					}
				}
			}
		}
		ERR_print_errors_fp(stderr);
		if (kctx[i])EVP_PKEY_CTX_free(kctx[i]); kctx[i] = NULL;
		if (params[i])EVP_PKEY_free(params[i]); params[i] = NULL;
		if (pctx[i])EVP_PKEY_CTX_free(pctx[i]); pctx[i] = NULL;
	}
	return rc;
}

int ecc_api_done()
{
	int i;
	for (i = 0; i < sizeof(_NIDS) / sizeof(_NIDS[0]); i++){
		if (kctx[i])EVP_PKEY_CTX_free(kctx[i]); kctx[i] = NULL;
		if (params[i])EVP_PKEY_free(params[i]); params[i] = NULL;
		if (pctx[i])EVP_PKEY_CTX_free(pctx[i]); pctx[i] = NULL;
	}
	return 0;
}

void * ecc_api_key_gen(ecc_pk_algorithm pk_alg, ecc_sym_algorithm sym_alg)
{
	EVP_PKEY * key = NULL;
	if (!EVP_PKEY_keygen(kctx[pk_alg], &key)){
		ERR_print_errors_fp(stderr);
	}
	return (void*)key;
}

void * ecc_api_key_init(ecc_pk_algorithm pk_alg, ecc_sym_algorithm sym_alg, const char* pkey)
{
	EVP_PKEY * key = NULL;
	BIGNUM * bn = BN_new();
	if (BN_bin2bn(pkey, 32, bn)){
		EC_KEY * eckey = EC_KEY_new_by_curve_name(_NIDS[pk_alg]);
		if (eckey){
			if (EC_KEY_set_private_key(eckey, bn)){
				EC_POINT * point;
				const EC_GROUP * group;
				group = EC_KEY_get0_group(eckey);
				point = EC_POINT_new(group);
				if (EC_POINT_mul(group, point, bn, NULL, NULL, NULL)){
					EC_KEY_set_public_key(eckey, point);
					key = EVP_PKEY_new();
					EVP_PKEY_set1_EC_KEY(key, eckey);
					eckey = NULL;
				}
				EC_POINT_free(point);
			}
			if (eckey){
				EC_KEY_free(eckey);
			}
		}
	}
	BN_free(bn);
	return (void*)key;
}

void   ecc_api_key_free(void* key)
{
	EVP_PKEY_free((EVP_PKEY*)key);
}

int ecc_api_key_private(void* key, char* buf)
{
	int len = -1;
	const EC_KEY   * eckey;
	eckey = EVP_PKEY_get1_EC_KEY((EVP_PKEY*)key);
	if(eckey){
		const BIGNUM   * ecbn;
		ecbn = EC_KEY_get0_private_key(eckey);
		if(ecbn){
			len = BN_num_bytes(ecbn);
			if(buf){
				BN_bn2bin(ecbn, (unsigned char*)buf);
			}
		}
	}
	return len;
}
int    ecc_api_key_public(void* key, char * px, char * py)
{
	const EC_KEY   * eckey;
	const EC_GROUP * ecgroup;
	const EC_POINT * ecpoint;
	BIGNUM x, y;
	int bcount = -1;
		
	eckey = EVP_PKEY_get1_EC_KEY(key);
	ecgroup = EC_KEY_get0_group(eckey);
	ecpoint = EC_KEY_get0_public_key(eckey);

	//fill public key data
	BN_init(&x); BN_init(&y);
	if(EC_POINT_get_affine_coordinates_GFp(ecgroup, ecpoint, &x, &y, NULL)){
		bcount = BN_num_bytes(&x);
		if(px){
			BN_bn2bin(&x, (unsigned char*)px);
		}

		bcount = BN_num_bytes(&y);
		if(py){
			BN_bn2bin(&x, (unsigned char*)py);
		}
	}
	BN_clear_free(&x); BN_clear_free(&y);
	return bcount;
}

static int _pass_cb(char *buf, int size, int rwflag, void *u)
{
	fprintf(stderr, "Ask for a pass phrase");
	return 0;
}
int    ecc_api_key_private_save(void* key, const char* path, ecc_format format)
{
	int rc = -1;
	FILE * f = fopen(path, "wb");
	if(f){
		if (format == ecc_pem){
			rc = PEM_write_PKCS8PrivateKey(f, key, NULL, NULL, 0, _pass_cb, NULL) ? 0 : -1;
			if (rc < 0){
				ERR_print_errors_fp(stderr);
			}
		}
		else{
			const EC_KEY   * eckey;
			eckey = EVP_PKEY_get1_EC_KEY((EVP_PKEY*)key);
			if (eckey){
				const BIGNUM   * ecbn;
				ecbn = EC_KEY_get0_private_key(eckey);
				if (ecbn){
					char * buf = NULL;
					int len = BN_num_bytes(ecbn);
					if (format == ecc_bin){
						buf = (char *)OPENSSL_malloc(len);
						BN_bn2bin(ecbn, buf);
						rc = 0;
					}
					else if (format == ecc_hex){
						buf = BN_bn2hex(ecbn);
						len = strlen(buf);
						rc = 0;
					}
					if (buf){
						rc = (len == fwrite(buf, 1, len, f)) ? 0 : -1;
						OPENSSL_free(buf);
					}
				}
			}
		}
		fclose(f);
		if (rc < 0){
			ERR_print_errors_fp(stderr);
			remove(path);
			rc = -1;
		}
	}
	else{
		perror(path);
	}
	return rc;
}

void * ecc_api_key_private_load(const char* path)
{
	EVP_PKEY * key = NULL;
	FILE * f = fopen(path, "rb");
	if (f){
		key = PEM_read_PrivateKey(f, NULL, NULL, NULL);
		if (key == NULL){
			BIGNUM * bn = NULL;
			fseek(f, 0, SEEK_END);
			int len = ftell(f);
			fseek(f, 0, SEEK_SET);
			char * buf = OPENSSL_malloc(len + 1);
			if (len == fread(buf, 1, len, f)){
				buf[len] = 0;
				// try hex first
				if (len != BN_hex2bn(&bn, buf)){
					if (bn){
						BN_free(bn); bn = NULL;
					}
					bn = BN_bin2bn(buf, len, NULL);
				}
			}
			OPENSSL_free(buf);
			if (bn){
				EC_KEY * eckey = EC_KEY_new_by_curve_name(_NIDS[0]);
				if (eckey){
					if (EC_KEY_set_private_key(eckey, bn)){
						EC_POINT * point;
						const EC_GROUP * group;
						group = EC_KEY_get0_group(eckey);
						point = EC_POINT_new(group);
						if (EC_POINT_mul(group, point, bn, NULL, NULL, NULL)){
							EC_KEY_set_public_key(eckey, point);
							key = EVP_PKEY_new();
							EVP_PKEY_assign_EC_KEY(key, eckey);
							eckey = NULL;
						}
						EC_POINT_free(point);
					}
					if (eckey)EC_KEY_free(eckey);
				}
				BN_free(bn);
			}
		}
		fclose(f);
	}
	return key;
}

int ecc_api_key_public_save(void* key , const char* path, ecc_format format)
{
	int rc = -1;
	FILE * f = fopen(path, "wb");
	if (f){
		const EC_KEY   * eckey;
		eckey = EVP_PKEY_get1_EC_KEY((EVP_PKEY*)key);
		if (eckey){
			if (format == ecc_pem){
				rc = PEM_write_EC_PUBKEY(f, eckey) ? 0 : -1;
			}
			else{
				size_t len;
				char * buf = NULL;
				const EC_POINT * point = EC_KEY_get0_public_key(eckey);
				const EC_GROUP * group = EC_KEY_get0_group(eckey);

				if (format == ecc_hex){
					buf = EC_POINT_point2hex(group, point, POINT_CONVERSION_UNCOMPRESSED, NULL);
					len = strlen(buf);
				}
				else if (format == ecc_bin){
					len = EC_POINT_point2oct(group, point, POINT_CONVERSION_UNCOMPRESSED, NULL, 0, NULL);
					if (len > 0){
						buf = OPENSSL_malloc(len + 1);
						if (len != EC_POINT_point2oct(group, point, POINT_CONVERSION_UNCOMPRESSED, buf, len, NULL)){
							OPENSSL_free(buf); buf = NULL;
						}
					}
				}
				if (buf){
					if (len == fwrite(buf, 1, len, f)){
						rc = 0;
					}
					OPENSSL_free(buf); buf = NULL;
				}
			}
		}
		fclose(f);
		if (rc < 0){
			ERR_print_errors_fp(stderr);
			remove(path);
		}
	}
	else{
		perror(path);
	}
	return rc;
}

void * ecc_api_key_public_load(const char* path, ecc_pk_algorithm pk_alg)
{
	EVP_PKEY * key = NULL;
	FILE * f = fopen(path, "rb");
	if (f){
		key = PEM_read_PUBKEY(f, &key, NULL, NULL);
		if (key == NULL){
			fseek(f, 0, SEEK_END);
			int len = ftell(f);
			fseek(f, 0, SEEK_SET);
			char * buf = OPENSSL_malloc(len + 1);
			if (len == fread(buf, 1, len, f)){
				buf[len] = 0;

				EC_KEY * eckey = EC_KEY_new_by_curve_name(_NIDS[pk_alg]);
				if (eckey){
					// try load hex
					EC_POINT * point = NULL;
					const EC_GROUP * group = EC_KEY_get0_group(eckey);
					// try hex first
					point = EC_POINT_hex2point(group, buf, NULL, NULL);
					if (point == NULL){
						// try oct
						point = EC_POINT_oct2point(group, NULL, buf, len, NULL);
					}
					if (point){
						EC_KEY_set_public_key(eckey, point);
						EC_POINT_free(point);
						key = EVP_PKEY_new();
						EVP_PKEY_assign_EC_KEY(key, eckey);
						eckey = NULL;
					}
					else{
						EC_KEY_free(eckey);
					}
				}
			}
			OPENSSL_free(buf);
		}
		fclose(f);
	}
	return key;
}

int sha256_calculate(char* hash, const char * ptr, int len)
{
	SHA256_CTX ctx;
	SHA256_Init(&ctx);
	SHA256_Update(&ctx, ptr, len);
	SHA256_Final((unsigned char*)hash, &ctx);
	return 0;
}

int    ecc_sign(void * key, const char * data, int length, char ** psig, int maxsize)
{
	EC_KEY   * eckey;
	unsigned char *sig = (unsigned char *)*psig; 

	if (65 <= maxsize){
		eckey = EC_KEY_dup(EVP_PKEY_get1_EC_KEY(key));
		if(eckey){
			unsigned char hash[32];
			ECDSA_SIG * ecdsa;
			SHA256_CTX ctx;
			BIGNUM *kinv = NULL;
			BIGNUM *rp = NULL;

			SHA256_Init(&ctx);
			SHA256_Update(&ctx, data, length);
			SHA256_Final(hash, &ctx);

			if (ECDSA_sign_setup(eckey, NULL, &kinv, &rp)){
				ecdsa = ECDSA_do_sign_ex(hash, 32, kinv, rp, eckey);
				EC_KEY_free(eckey);
				BN_clear_free(kinv);
				BN_clear_free(rp);
				if (ecdsa){
					int bcount;
					*(sig++) = 0; // x_coordinate_only
					bcount = BN_num_bytes(ecdsa->r);
					BN_bn2bin(ecdsa->r, sig);
					sig += bcount;
					bcount = BN_num_bytes(ecdsa->s);
					BN_bn2bin(ecdsa->s, sig);
					sig += bcount;
					ECDSA_SIG_free(ecdsa);
					*psig = (char*)sig;
					return 0;
				}
			}
		}
	}
	return -1;
}
