#include <openssl/pem.h>
#include <openssl/evp.h>
#include <openssl/bio.h>
#include <openssl/err.h>
#include <iostream>
#include <iomanip>
#include <fstream>
#include <cstring>
#include <algorithm>


using namespace std;


int generate_key()
{
	ofstream fout;
	BIO* bio;

	// generate key pair
	EVP_PKEY* key = NULL;
	EVP_PKEY_CTX *pctx = EVP_PKEY_CTX_new_id(EVP_PKEY_ED25519, NULL);
	if(EVP_PKEY_keygen_init(pctx) != 1)
	{
		cout << "Error: EVP_PKEY_keygen_init" << endl;
		cout << ERR_error_string(ERR_get_error(),NULL) << endl;
		return 0;
	}
	if(EVP_PKEY_keygen(pctx, &key) != 1)
	{
		cout << "Error: EVP_PKEY_keygen" << endl;
		cout << ERR_error_string(ERR_get_error(),NULL) << endl;
		return 0;
	}
	EVP_PKEY_CTX_free(pctx);
	
	// save private key
	bio = BIO_new( BIO_s_mem() );
	if(PEM_write_bio_PrivateKey(bio, key, NULL, NULL, 0, NULL, NULL) != 1)
	{
		cout << "Error: PEM_write_bio_PrivateKey" << endl;
		cout << ERR_error_string(ERR_get_error(),NULL) << endl;
		return 0;
	}
	unsigned char private_key[120] = { 0 };
	BIO_read(bio, private_key, 120);
	BIO_free(bio);
	cout << private_key << endl;
	fout.open("private_key");
	fout << private_key;
	fout.close();	

	// save public key
	bio = BIO_new( BIO_s_mem() );
	if(PEM_write_bio_PUBKEY(bio, key) != 1)
	{
		cout << "Error: PEM_write_bio_PUBKEY" << endl;
		cout << ERR_error_string(ERR_get_error(),NULL) << endl;
		return 0;
	}
	EVP_PKEY_free(key);
	unsigned char public_key[120] = { 0 };
	BIO_read(bio, public_key, 120);
	BIO_free(bio);
	cout << public_key << endl;
	fout.open("public_key");
	fout << public_key;
	fout.close();

	return 1;
}



int sign(char* signee_file, char* key_file)
{
	ifstream fin;
	ofstream fout;	
	BIO* bio;
	EVP_PKEY* pkey = NULL;
	int signee_size = 0, key_size = 0;
	
	// open signee_file
	fin.open(signee_file);
	if (fin.fail()) {
		cout << "Error: Can't open signee file." << endl;
		return 0;
	}
	fin.seekg(0, ios::end);
	signee_size = fin.tellg();
	fin.seekg(0, ios::beg);
	unsigned char* signee = new unsigned char[signee_size];
	fin.read((char*)signee, signee_size);
	fin.close();

	// open private key
	fin.open(key_file);
	if (fin.fail()) {
		cout << "Error: Can't open key file." << endl;
		return 0;
	}
	fin.seekg(0, ios::end);
	key_size = fin.tellg();
	fin.seekg(0, ios::beg);
	unsigned char* key = new unsigned char[key_size];
	fin.read((char*)key, key_size);
	fin.close();
	cout << key << endl;
	bio = BIO_new( BIO_s_mem() );
	BIO_write(bio, key, key_size);
	if(PEM_read_bio_PrivateKey(bio, &pkey, NULL, NULL) == NULL)
	{	
		cout << "Error: PEM_read_bio_PrivateKey" << endl;
		cout << ERR_error_string(ERR_get_error(),NULL) << endl;
		return 0;
	}
	BIO_free(bio);

	// File Signing, Make Signature
	unsigned char sig[65] = { 0 };
	size_t sig_size = 64;
	EVP_MD_CTX* ctx = EVP_MD_CTX_new();
	if(EVP_DigestSignInit(ctx, NULL, NULL, NULL, pkey) != 1)
	{
		cout << "Error: EVP_DigestSignInit" << endl;
		cout << ERR_error_string(ERR_get_error(),NULL) << endl;
		return 0;
	}
	if(EVP_DigestSign(ctx, sig, &sig_size, signee, signee_size) != 1)
	{
		cout << "Error: EVP_DigestSign" << endl;
		cout << ERR_error_string(ERR_get_error(),NULL) << endl;
		return 0;
	}
	EVP_MD_CTX_free(ctx);
	EVP_PKEY_free(pkey);

	fout.open("./signature");
	cout << hex << setfill('0');
	fout << hex << setfill('0');
	for(int i = 0; i < sig_size; i++)
	{
		cout << setw(2) << (int)sig[i];
		fout << setw(2) << (int)sig[i];
	}
	cout << endl;
	fout.close();

	return 1;
}


/*
int verify(char* signee_file, char* sig_file, char* key_file)
{
	ifstream fin;
	BIO* bio;
	EVP_PKEY* pkey = NULL;
	int signee_size = 0, key_size = 0, sig_size = 0;
	
	// open signee_file
	fin.open(signee_file);
	fin.seekg(0, ios::end);
	signee_size = fin.tellg();
	fin.seekg(0, ios::beg);
	unsigned char* signee = new unsigned char[signee_size];
	fin.read((char*)signee, signee_size);
	fin.close();

	// open public key
	fin.open(key_file);
	fin.seekg(0, ios::end);
	key_size = fin.tellg();
	fin.seekg(0, ios::beg);
	unsigned char* key = new unsigned char[key_size];
	fin.read((char*)key, key_size);
	fin.close();
	cout << key << endl;
	bio = BIO_new( BIO_s_mem() );
	BIO_write(bio, key, key_size);
	if(PEM_read_bio_PUBKEY(bio, &pkey, NULL, NULL) == NULL)
	{	
		cout << "Error: " << ERR_error_string(ERR_get_error(),NULL) << endl;
		return 0;
	}
	BIO_free(bio);

	fin.open(sig_file);
	fin.seekg(0, ios::end);
	sig_size = fin.tellg();
	fin.seekg(0, ios::beg);
	unsigned char* sig_char = new unsigned char[sig_size];
	fin.read((char*)sig_char, sig_size);
	fin.close();
	unsigned char sig[65] = { 0 };
	char temp[3] = { 0 };
	for(int i = 0; i < sig_size; i+=2)
	{
		temp[0] = sig_char[i];
		temp[1] = sig_char[i+1];
		sig[i/2] = (unsigned char) strtol(temp, NULL, 16);
	}
	sig_size /= 2;

	EVP_MD_CTX* ctx = EVP_MD_CTX_new();
	if(EVP_DigestVerifyInit(ctx, NULL, NULL, NULL, pkey) != 1)
	{
		cout << "Error: " << ERR_error_string(ERR_get_error(),NULL) << endl;
		return 0;
	}
	if(EVP_DigestVerify(ctx, sig, sig_size, signee, signee_size) != 1)
	{
		cout << "Error: " << ERR_error_string(ERR_get_error(),NULL) << endl;
		return 0;
	}
	EVP_PKEY_free(pkey);
	
	cout << "Success" << endl;
	return 1;
}
*/


bool cmdOptionExists(char** begin, char** end, const string& option)
{
	return find(begin, end, option) != end;
}



char* getCmdOption(char** begin, char** end, const string& option)
{
	char** itr = find(begin, end, option);
	if (itr != end && ++itr != end)
		return *itr;
	return 0;
}



int main(int argc, char* argv[])
{
	if (cmdOptionExists(argv, argv+argc, "-h")\
		|| cmdOptionExists(argv, argv+argc, "--help"))
	{
		cout << "Usage" << endl;
		cout << "Generate key: " << argv[0] << " --genkey" << endl;
		cout << "Signing file: " << argv[0] << " -f signee_path -k key_path" << endl;
		return 0;
	}

	if (cmdOptionExists(argv, argv+argc, "--genkey"))
	{
		generate_key();
		return 0;
	}

	if (cmdOptionExists(argv, argv+argc, "-f")\
		&& cmdOptionExists(argv, argv+argc, "-k"))
	{
		char* signee_path = getCmdOption(argv, argv+argc, "-f");
		char* key_path = getCmdOption(argv, argv+argc, "-k");
		if (sign(signee_path, key_path) != 1)
			cout << "Error occured" << endl;
		return 0;
	}
	else
	{
		cout << "Usage" << endl;
		cout << "Generate key: " << argv[0] << " --genkey" << endl;
		cout << "Signing file: " << argv[0] << " -f signee_path -k key_path" << endl;
	}
//	generate_key();
//	sign(argv[1], argv[2]);
//	verify(argv[1], argv[2], argv[3]);
	return 0;
}