/*
	This file is part of Bolixo.

	Bolixo is free software: you can redistribute it and/or modify
	it under the terms of the GNU General Public License as published by
	the Free Software Foundation, either version 3 of the License, or
	(at your option) any later version.

	Bolixo is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU General Public License for more details.

	You should have received a copy of the GNU General Public License
	along with Bolixo.  If not, see <https://www.gnu.org/licenses/>.
*/
#include <stdio.h>
#include <string.h>
#include <tlmplib.h>
#include <memory>
#include <string>

using namespace std;
using std::unique_ptr;

#include <openssl/bn.h>
#include <openssl/rsa.h>
#include <openssl/pem.h>
#include <openssl/bio.h>
#include <openssl/x509.h>

#include <cassert>
#define ASSERT assert

using BN_ptr = std::unique_ptr<BIGNUM, decltype(&::BN_free)>;
//using RSA_ptr = std::unique_ptr<RSA, decltype(&::RSA_free)>;
using EVP_KEY_ptr = std::unique_ptr<EVP_PKEY, decltype(&::EVP_PKEY_free)>;
using BIO_FILE_ptr = std::unique_ptr<BIO, decltype(&::BIO_free)>;


static int sign(
	const char *msg,
	EVP_PKEY *key,
	unsigned char *&sig,
	size_t &sig_len)
{
	int ret = -1;
	sig = NULL;
	sig_len = 0;
	/* Create the Message Digest Context */
	EVP_MD_CTX *mdctx = EVP_MD_CTX_create();
	if (mdctx != NULL){
		/* Initialise the DigestSign operation - SHA-256 has been selected as the message digest function in this example */
		if(EVP_DigestSignInit(mdctx, NULL, EVP_sha256(), NULL, key) == 1){
			/* Call update with the message */
			if(EVP_DigestSignUpdate(mdctx, msg, strlen(msg)) == 1){
 
				// Finalise the DigestSign operation
				// First call EVP_DigestSignFinal with a NULL sig parameter to obtain the length of the
				// signature. Length is returned in slen */
				if(EVP_DigestSignFinal(mdctx, NULL, &sig_len) == 1){
					//printf ("sig_len=%lu\n",sig_len);
					// Allocate memory for the signature based on size in slen
					sig = (unsigned char*)OPENSSL_malloc(sizeof(unsigned char) * sig_len);
					if (sig != NULL){
						// Obtain the signature
						if(EVP_DigestSignFinal(mdctx, sig, &sig_len) == 1){
							/* Success */
							#if 0
								for (unsigned i=0; i<sig_len; i++){
									printf ("%02x ",sig[i]);
								}
								printf ("\n");
							#endif
							ret = 1;
						}else{
							OPENSSL_free(sig);
							sig = NULL;
							// error
						}
					} 
				}
			}
		}
		EVP_MD_CTX_destroy(mdctx);
	}
	return ret;
}

static int verify( 
	PARAM_STRING msg,
	EVP_PKEY *key,
	unsigned char *sig,
	size_t sig_len)
{
	int ret = -1;
	/* Initialize `key` with a public key */
	EVP_MD_CTX *mdctx = EVP_MD_CTX_create();
	if (mdctx != NULL){
		if(EVP_DigestVerifyInit(mdctx, NULL, EVP_sha256(), NULL, key) == 1){
			// Initialize `key` with a public key
			if(EVP_DigestVerifyUpdate(mdctx, msg.ptr, strlen(msg.ptr)) == 1){
				if(EVP_DigestVerifyFinal(mdctx, sig, sig_len) == 1){
					// Success
					ret = 0;
				}else{
					// Failure
				}
			}
		}
		EVP_MD_CTX_destroy(mdctx);
	}
	return ret;
}
#ifndef OPENSSL_VERSION_MAJOR
	#define OLDRSA
#endif
#ifdef OLDRSA
static RSA *generate(unsigned nbbits, string &private_key, string &public_key, PARAM_STRING passphrase, bool write_files)
{
	private_key.clear();
	public_key.clear();
	int rc;

	RSA *rsa = RSA_new();
	BN_ptr bn(BN_new(), ::BN_free);


	rc = BN_set_word(bn.get(), RSA_F4);
	ASSERT(rc == 1);

	// Generate key
	rc = RSA_generate_key_ex(rsa, nbbits, bn.get(), NULL);
	ASSERT(rc == 1);

	// Convert RSA to PKEY
	EVP_KEY_ptr pkey(EVP_PKEY_new(), ::EVP_PKEY_free);
	rc = EVP_PKEY_set1_RSA(pkey.get(), rsa);
	ASSERT(rc == 1);

	{
		BIO *bio = BIO_new(BIO_s_mem());
		PEM_write_bio_PrivateKey(bio, pkey.get(),EVP_aes_192_cbc(), (unsigned char*)passphrase.ptr,strlen(passphrase.ptr), 0, 0);
		char line[1000];
		while(BIO_gets(bio,line,sizeof(line))>0){
			//printf ("line=%s\n",line);
			private_key += line;
		}
		BIO_free_all(bio);
	}
	{
		BIO *bio = BIO_new(BIO_s_mem());
		PEM_write_bio_PUBKEY(bio, pkey.get());
		char line[1000];
		while(BIO_gets(bio,line,sizeof(line))>0){
			//printf ("line=%s\n",line);
			public_key += line;
		}
		BIO_free_all(bio);
	}
	if (write_files){
		BIO_FILE_ptr pem1(BIO_new_file("/tmp/rsa-public-1.pem", "w"), ::BIO_free);
		BIO_FILE_ptr pem2(BIO_new_file("/tmp/rsa-public-2.pem", "w"), ::BIO_free);
		BIO_FILE_ptr pem3(BIO_new_file("/tmp/rsa-private-1.pem", "w"), ::BIO_free);
		BIO_FILE_ptr pem4(BIO_new_file("/tmp/rsa-private-2.pem", "w"), ::BIO_free);
		BIO_FILE_ptr pem5(BIO_new_file("/tmp/rsa-private-3.pem", "w"), ::BIO_free);
		BIO_FILE_ptr der1(BIO_new_file("/tmp/rsa-public.der", "w"), ::BIO_free);
		BIO_FILE_ptr der2(BIO_new_file("/tmp/rsa-private.der", "w"), ::BIO_free);

		// Write public key in ASN.1/DER
		rc = i2d_RSAPublicKey_bio(der1.get(), rsa);
		ASSERT(rc == 1);

		// Write public key in PKCS PEM
		rc = PEM_write_bio_RSAPublicKey(pem1.get(), rsa);
		ASSERT(rc == 1);       

		// Write public key in Traditional PEM
		rc = PEM_write_bio_PUBKEY(pem2.get(), pkey.get());
		ASSERT(rc == 1);
		//////////

		// Write private key in ASN.1/DER
		rc = i2d_RSAPrivateKey_bio(der2.get(), rsa);
		ASSERT(rc == 1);

		// Write private key in PKCS PEM.
		rc = PEM_write_bio_PrivateKey(pem3.get(), pkey.get(), NULL, NULL, 0, NULL, NULL);
		ASSERT(rc == 1);

		// Write private key in PKCS PEM
		rc = PEM_write_bio_PKCS8PrivateKey(pem4.get(), pkey.get(), NULL, NULL, 0, NULL, NULL);
		ASSERT(rc == 1);

		// Write private key in Traditional PEM
		rc = PEM_write_bio_RSAPrivateKey(pem5.get(), rsa, NULL, NULL, 0, NULL, NULL);
		ASSERT(rc == 1);
	}
	return rsa;
}
#else
static void get_private_key (EVP_PKEY *pkey, PARAM_STRING passphrase, string &private_key)
{
	BIO *bio = BIO_new(BIO_s_mem());
	PEM_write_bio_PrivateKey(bio, pkey,EVP_aes_192_cbc(), (unsigned char*)passphrase.ptr,strlen(passphrase.ptr), 0, 0);
	char line[1000];
	while(BIO_gets(bio,line,sizeof(line))>0){
		//printf ("line=%s\n",line);
		private_key += line;
	}
	BIO_free_all(bio);
}
static void get_public_key (EVP_PKEY *pkey, string &public_key)
{
	BIO *bio = BIO_new(BIO_s_mem());
	PEM_write_bio_PUBKEY(bio, pkey);
	char line[1000];
	while(BIO_gets(bio,line,sizeof(line))>0){
		//printf ("line=%s\n",line);
		public_key += line;
	}
	BIO_free_all(bio);
}
static void generate_new(unsigned nbbits, string &private_key, string &public_key, PARAM_STRING passphrase, bool write_files)
{
	private_key.clear();
	public_key.clear();
	ENGINE *e = nullptr;
	EVP_PKEY_CTX *ctx = EVP_PKEY_CTX_new_id(EVP_PKEY_RSA,e);
	if (ctx == nullptr){
		fprintf (stderr,"ctx=%p\n",ctx);
		exit (-1);
	}

	int rc;
	#if 0
		// Does not work for RSA
		rc = EVP_PKEY_paramgen_init (ctx);
		printf ("paramgen_init rc=%d\n",rc);

		const OSSL_PARAM *params = EVP_PKEY_CTX_gettable_params(ctx);
		printf ("params=%p\n",params);
	#endif
	// Generate key
	rc = EVP_PKEY_keygen_init (ctx);
	printf ("keygen_init rc=%d\n",rc);
	rc = EVP_PKEY_CTX_set_rsa_keygen_bits(ctx, nbbits);
	printf ("keygen_bits rc=%d\n",rc);

	EVP_PKEY *retkey=nullptr;
	rc = EVP_PKEY_generate(ctx,&retkey);
	printf ("generate rc=%d\n",rc);
	get_private_key(retkey,passphrase,private_key);
	get_public_key(retkey,public_key);
	//printf ("retkey.size=%d bits=%d\n",EVP_PKEY_get_size(retkey),EVP_PKEY_get_bits(retkey));
	//printf ("private_key=%s\npublic_key=%s\n",private_key.c_str(),public_key.c_str());
	EVP_PKEY_CTX_free(ctx);
	EVP_PKEY_free(retkey);
}
#endif

static EVP_PKEY *testsign_load_private (PARAM_STRING p, PARAM_STRING passphrase)
{
	BIO *bio = BIO_new_mem_buf((void*)p.ptr, strlen(p.ptr));
	EVP_PKEY *pkey = PEM_read_bio_PrivateKey(bio, NULL,NULL,(void*)passphrase.ptr);
	BIO_free_all (bio);
	return pkey;
}
static EVP_PKEY *testsign_load_public (PARAM_STRING p)
{
	BIO *bio = BIO_new_mem_buf((void*)p.ptr, strlen(p.ptr));
	EVP_PKEY *pkey = PEM_read_bio_PUBKEY(bio, NULL,NULL,NULL);
	BIO_free_all (bio);
	return pkey;
}
<mod>
int main (int argc, char *argv[])
{
	glocal int ret = -1;
	glocal unsigned nbverify = 1;
	glocal unsigned nbgen = 1;
	glocal unsigned nbbits = 2048;
	glocal const char *message =
		"This is a message that goes like this\n"
		"This is a message that goes like this\n"
		"This is a message that goes like this\n"
		"This is a message that goes like this\n"
		"This is a message that goes like this\n"
		"This is a message that goes like this\n"
		"This is a message that goes like this\n"
		"This is a message that goes like this\n"
		"This is a message that goes like this\n"
		"This is a message that goes like this\n"
		"This is a message that goes like this\n"
		"This is a message that goes like this\n"
		"This is a message that goes like this\n"
		;
	glocal.ret = <call tlmpprogram>(argc,argv);
	<f init>
		setproginfo ("testsign",VERSION,"...");
		setarg (' ',"nbverify","Number of sign/verifies",glocal.nbverify,false);
		setarg (' ',"nbgen","Number of key generation",glocal.nbgen,false);
		setarg (' ',"nbbits","Size in bits of the generated key",glocal.nbbits,false);
	</f>
	<f main_noarg>
		string modified_message = string(glocal.message) + "something";
		const char *passphrase = "this is secret";
		int ret = -1;
		#ifdef OLDRSA
		RSA *rsa = NULL;
		#endif
		string private_key,public_key;
		for (unsigned i=0; i<glocal.nbgen; i++){
			#ifdef OLDRSA
				rsa = generate(glocal.nbbits,private_key,public_key,passphrase,false);
				if (i <glocal.nbgen-1) RSA_free (rsa);
			#else
				generate_new(glocal.nbbits,private_key,public_key,passphrase,false);
			#endif
		}
		printf ("private=\n%s\npublic\n%s\n",private_key.c_str(),public_key.c_str());
			
		{
			{
				EVP_PKEY *privkey = testsign_load_private(private_key,passphrase); 
				if (privkey != NULL){
					#ifndef OLDRSA
						printf ("privkey.size=%d bits=%d\n",EVP_PKEY_get_size(privkey),EVP_PKEY_get_bits(privkey));
					#endif
					EVP_PKEY *pubkey = testsign_load_public(public_key); 
					if (pubkey != NULL){
						#ifndef OLDRSA
							printf ("pubkey.size=%d bits=%d\n",EVP_PKEY_get_size(pubkey),EVP_PKEY_get_bits(pubkey));
						#endif
						unsigned char *sig;
						size_t sig_len;
						if (sign (glocal.message,privkey,sig,sig_len)!=-1){
							printf ("sign1 %p %zu\n",sig,sig_len);
							if (verify (glocal.message,pubkey,sig,sig_len)!=-1){
								printf ("Verify1 ok\n");
							}else{
								printf ("Verify1 pas ok\n");
							}
							if (verify (modified_message,pubkey,sig,sig_len)!=-1){
								printf ("Verify2 ok\n");
							}else{
								printf ("Verify2 pas ok\n");
							}
							OPENSSL_free (sig);
						}
						EVP_PKEY_free (pubkey);
					}
					EVP_PKEY_free (privkey);
				}
			}
			#ifdef OLDRSA
			EVP_PKEY *pkey = EVP_PKEY_new(); 
			if (pkey != NULL){
				EVP_PKEY_assign_RSA(pkey, rsa); 
				for (unsigned i=0; i<glocal.nbverify; i++){
					unsigned char *sig;
					size_t sig_len;
					if (sign (glocal.message,pkey,sig,sig_len)!=-1){
						if (verify (glocal.message,pkey,sig,sig_len)!=-1){
							printf ("Verify3 ok\n");
						}else{
							printf ("Verify3 pas ok\n");
						}
						if (verify (modified_message,pkey,sig,sig_len)!=-1){
							printf ("Verify4 ok\n");
						}else{
							printf ("Verify4 pas ok\n");
						}
						OPENSSL_free (sig);
						ret = 0;
					}
				}
				EVP_PKEY_free (pkey);
			}else{
				RSA_free (rsa);
			}
			#endif
		}
		return ret;
	</f>
	</call>
	return glocal.ret;
}
</mod>