#include <assert.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

// openssl
#include <openssl/err.h>
#include <openssl/bn.h>
#include <openssl/sha.h>
#include <openssl/rsa.h>
#include <openssl/x509.h>
#include <openssl/aes.h>
#include <openssl/pem.h>

#include "nfastapp.h"
#include "nfkm.h"
#include "rqcard-applic.h"
#include "rqcard-fips.h"

//#include "picky-upcalls.h"

#include "my_hsm_bignum.h"
#include "my_hsm_alloc.h"
#include "my_hsm_setup.h"

#define PRIV_KEY_FILE		"./test_key/test-rsa-privkey2048.der"
#define PUB_KEY_FILE		"./test_key/test-rsa-pubkey2048.der"

#define DATA_LEN		256	// bytes

// RSA private key data
typedef struct
{
	struct NFast_Bignum *p;
	struct NFast_Bignum *q;
	struct NFast_Bignum *dmp1;
	struct NFast_Bignum *dmq1;
	struct NFast_Bignum *iqmp;
	struct NFast_Bignum *e;
}
RSAPrivateKeyDataBn;

// RSA public key data
typedef struct
{
	struct NFast_Bignum *e;
	struct NFast_Bignum *n;
}
RSAPublicKeyDataBn;

// global variable
NFast_AppHandle handle;
NFastApp_Connection nc;
NFKM_WorldInfo *pWorld = NULL;
RQCard card;
RQCard_FIPS fips;
NFKM_ModuleInfo *pModuleinfo = NULL;
M_KeyID ltid = 0;	// the cardset loaded into the module

const NFKM_KeyIdent priv_keyident = { (char*)"simple", (char*)"test-rsa-privkey2048" };
const NFKM_KeyIdent pub_keyident = { (char*)"simple", (char*)"test-rsa-pubkey2048" };

unsigned char save_enc[DATA_LEN];

// function
int createRSAPSSCert( NFKM_KeyIdent priv_ident, NFKM_KeyIdent pub_ident );
void PrintArray( char *pStr, const unsigned char *pData, int length );

int createRSAPSSCert( NFKM_KeyIdent priv_ident, NFKM_KeyIdent pub_ident )
{
	int i;
	int result = Status_OK;
	M_ByteBlock *blobptr = NULL;
	M_KeyID priv_keyid, pub_keyid;
	NFKM_Key *keyinfo = NULL;
	M_Command cmd;
	M_Reply reply;

	priv_keyid = pub_keyid = 0;
	memset( &cmd, 0, sizeof( cmd ) );
	memset( &reply, 0, sizeof( reply ) );

	// find priv key
	result = NFKM_findkey( handle, priv_ident, &keyinfo, NULL );
	if ( result != Status_OK )
	{	
		printf( "error(%d) : NFKM_findkey(priv)\n", result );
		return result;
	}
    if ( keyinfo == NULL )
    {
        printf( "error : not found key(priv)\n" );
        return Status_UnknownKey;
    }
	
	// load priv key blob
	blobptr = &(keyinfo->privblob);
	result = NFKM_cmd_loadblob( handle, nc, 
		pModuleinfo->module, blobptr, ltid, &priv_keyid, "loading priv key blob", NULL );
	if ( result != Status_OK )
	{	
		printf( "error(%d) : NFKM_cmd_loadblob(priv)\n", result );
		return result;
	}
	NFKM_freekey( handle, keyinfo, NULL );
	keyinfo = NULL;

#if 0
	// get priv key info
	cmd.cmd = Cmd_GetKeyInfo;
	cmd.args.getkeyinfo.key = priv_keyid;
	result = NFastApp_Transact( nc, NULL, &cmd, &reply, NULL );
	if ( result != Status_OK )
	{	
		printf( "error(%d) : FastApp_Transact(Cmd_GetKeyInfo)\n", result );
		return result;
	}
	printf( "priv key ID  : %08X\n", (unsigned int)priv_keyid );			
	printf( "priv keytype : %d\n", reply.reply.getkeyinfo.type );
	NFastApp_Free_Command( handle, NULL, NULL, &cmd );
	NFastApp_Free_Reply( handle, NULL, NULL, &reply );
	memset( &cmd, 0, sizeof( cmd ) );
	memset( &reply, 0, sizeof( reply ) );
#endif
		
	// find pub key
	result = NFKM_findkey( handle, pub_ident, &keyinfo, NULL );
	if ( result != Status_OK )
	{	
		printf( "error(%d) : NFKM_findkey(pub)\n", result );
		return result;
	}
    if ( keyinfo == NULL )
    {
        printf( "error : not found key(pub)\n" );
        return Status_UnknownKey;
    }
	
	// load pub key blob
	blobptr = &(keyinfo->pubblob); // pub dakedo privblob
	result = NFKM_cmd_loadblob( handle, nc, 
		pModuleinfo->module, blobptr, ltid, &pub_keyid, "loading pub key blob", NULL );
	if ( result != Status_OK )
	{	
		printf( "error(%d) : NFKM_cmd_loadblob(pub)\n", result );
		return result;
	}
	NFKM_freekey( handle, keyinfo, NULL );
	keyinfo = NULL;
	
#if 0
	// get priv key info
	cmd.cmd = Cmd_GetKeyInfo;
	cmd.args.getkeyinfo.key = pub_keyid;
	result = NFastApp_Transact( nc, NULL, &cmd, &reply, NULL );
	if ( result != Status_OK )
	{	
		printf( "error(%d) : FastApp_Transact(Cmd_GetKeyInfo)\n", result );
		return result;
	}
	printf( "pub key ID  : %08X\n", (unsigned int)pub_keyid );
	printf( "pub keytype : %d\n", reply.reply.getkeyinfo.type );
	NFastApp_Free_Command( handle, NULL, NULL, &cmd );
	NFastApp_Free_Reply( handle, NULL, NULL, &reply );
	memset( &cmd, 0, sizeof( cmd ) );
	memset( &reply, 0, sizeof( reply ) );
	keyinfo = NULL;
#endif

    // create cert
    {
        unsigned char *beforePtr, *middlePtr, *afterPtr;
		int beforeLen, middleLen, afterLen;
		unsigned char *signPtr, *decPtr;
		struct NFast_Bignum *sign_bn, *dec_bn;
		M_Mech verify_mech;
				
		beforePtr = middlePtr = afterPtr = NULL;
		beforeLen = middleLen = afterLen = 0;
		signPtr = decPtr = NULL;
		sign_bn = dec_bn = NULL;
		
        M_Hash32 hash32;
        M_Hash32 modHash32;

        unsigned int magic = 0xdeadbeef;
        SHA256( (const unsigned char*)&magic, 4, hash32.bytes );

	    // HSM は切り詰めないで署名してしまうので自前で加工する
	    memset( &modHash32, 0, sizeof( modHash32 ) );
        modHash32.bytes[2] = hash32.bytes[0] >> 7;
    	for( i = 3; i < SHA256_DIGEST_LENGTH; i++ )
        {
            modHash32.bytes[i] = (hash32.bytes[i-3] << 1) | (hash32.bytes[i-2] >> 7);
        }

		// sign data setting
		beforeLen = DATA_LEN - 11;
		beforePtr = (unsigned char*)malloc( beforeLen );
		for ( i = 0; i < beforeLen; i++ )
			beforePtr[i] = ~i;
		my_bin2bignum( &(sign_bn), handle, beforePtr, beforeLen );

        // sign transact
        cmd.cmd = Cmd_Sign;
        cmd.args.sign.flags = 0;
        cmd.args.sign.key = priv_keyid;
//        cmd.args.sign.mech = Mech_RSAhSHA256pPSS;
        cmd.args.sign.mech = Mech_RSAhSHA256pPKCS1;
        cmd.args.sign.plain.type = PlainTextType_Hash32;
        cmd.args.sign.plain.data.hash32.data = hash32;
//        cmd.args.sign.plain.data.hash32.data = modHash32;
//        RSA_verify_PKCS1_PSS
//        cmd.args.sign.plain.data.bignum.m = sign_bn;

		result = NFastApp_Transact( nc, NULL, &cmd, &reply, NULL );
		if ( result != Status_OK )
		{	
			printf( "error(%d) : FastApp_Transact(Cmd_Sign)\n", result );
			return 1;
		}
		result = reply.status;
		if ( result != Status_OK )
		{	
			printf( "error(%d) : reply.status(Cmd_Sign)\n", result );
			return 1;
		}

		printf( "RSA PSS sign ok\n" );

        // verify data setting
        middleLen = reply.reply.sign.sig.data.rsappkcs1.m->nbytes;
		middlePtr = (unsigned char*)malloc( middleLen );
		my_bignum2bin( middlePtr, middleLen, handle,
            reply.reply.sign.sig.data.rsappkcs1.m );
		my_bin2bignum( &dec_bn, handle, middlePtr, middleLen );		
        verify_mech = reply.reply.sign.sig.mech;
				
		NFastApp_Free_Command( handle, NULL, NULL, &cmd );
		NFastApp_Free_Reply( handle, NULL, NULL, &reply );

        memset( &cmd, 0, sizeof( cmd ) );
		memset( &reply, 0, sizeof( reply ) );

		// verify transact
        cmd.cmd = Cmd_Verify;
        cmd.args.verify.flags = 0;
		cmd.args.verify.key = pub_keyid;
//        cmd.args.verify.mech = Mech_RSAhSHA256pPSS;
        cmd.args.verify.mech = Mech_RSAhSHA256pPKCS1;
        cmd.args.verify.plain.type = PlainTextType_Hash32;
        cmd.args.verify.plain.data.hash32.data = hash32;
//        cmd.args.verify.plain.data.hash32.data = modHash32;
        cmd.args.verify.sig.mech = verify_mech;
        cmd.args.verify.sig.data.rsappkcs1.m = dec_bn;

//		cmd.args.decrypt.cipher.data.rsappkcs1.m = dec_bn;
//		cmd.args.decrypt.reply_type = PlainTextType_Bignum;
		result = NFastApp_Transact( nc, NULL, &cmd, &reply, NULL );
        if ( result != Status_OK )
		{	
			printf( "error(%d) : FastApp_Transact(Cmd_Verify)\n", result );
			return 1;
		}
        result = reply.status;
        if ( result != Status_OK && result != Status_VerifyFailed )
		{	
			printf( "error(%d) : reply.status(Cmd_Verify)\n", result );
			return 1;
		}
        // リファンレスは間違い。reply の status に検証結果が入る
        if ( result == Status_VerifyFailed )
        {
			printf( "error : Verify Failed(Cmd_Verify)\n" );
            return 1;
        }

		printf( "RSA PSS verify ok\n" );

        PrintArray( (char*)"HSM Sign", middlePtr, middleLen );

#if 1
        // OpenSSL でも検証開始
        {
            int ret = 0;

            unsigned char *privKeyBuf = NULL;
            unsigned char *pubKeyBuf = NULL;
            int privKeySize = 0;
            int pubKeySize = 0;

            RSA *rsa_privkey = NULL;
            RSA *rsa_pubkey  = NULL;

            FILE *pPrivKeyFile;
            FILE *pPubKeyFile;
            struct stat fileStat;

            // priv
            stat( PRIV_KEY_FILE, &fileStat );
            privKeySize = fileStat.st_size;
            privKeyBuf = (unsigned char*)malloc( privKeySize );
            pPrivKeyFile = fopen( PRIV_KEY_FILE, "rb" );
            if (pPrivKeyFile == NULL)
            {
                printf("file open error\n");
                return 1;
            }
            fread( privKeyBuf, privKeySize, 1, pPrivKeyFile );
            fclose(pPrivKeyFile);

            // pub
            stat( PUB_KEY_FILE, &fileStat );
            pubKeySize = fileStat.st_size;
            pubKeyBuf = (unsigned char*)malloc( pubKeySize );
            pPubKeyFile = fopen( PUB_KEY_FILE, "rb" );
            if (pPubKeyFile == NULL)
            {
                printf("file open error\n");
                return 1;
            }
            fread( pubKeyBuf, pubKeySize, 1, pPubKeyFile );
            fclose(pPubKeyFile);

            // DERフォーマットのRSA鍵を読み込み
            {
                const unsigned char *der_priv = privKeyBuf;
                const unsigned char *der_pub  = pubKeyBuf;

//                int priv_len = der_priv[ 8 ] | der_priv[ 9 ] << 8;	// KEY長を取り出し
//                int pub_len  = der_pub [ 8 ] | der_pub [ 9 ] << 8;  // 同上
//                der_priv += 0x10; // ヘッダ部分を除外してKEY実体を指定
//                der_pub  += 0x10; // 同上

                // コマンドラインのopensslが出力する秘密鍵は、PKCS#1 RSAPublicKeyフォーマットなので、この関数を使う。
                rsa_privkey = d2i_RSAPrivateKey( NULL, &der_priv, privKeySize );
                if( rsa_privkey == NULL )
                {
                    printf( "error : d2i_RSAPrivateKey\n" );
                    return 1;
                }
                // コマンドラインのopensslが出力する公開鍵は、SubjectPublicKeyInfo形式なので、この関数を使う。
                rsa_pubkey = d2i_RSA_PUBKEY( NULL, &der_pub, pubKeySize );
                if( rsa_pubkey == NULL )
                {
                    printf( "error : d2i_RSA_PUBKEY\n" );
                    return 1;
                }
            }

            PrintArray( (char*)"hash32", hash32.bytes, sizeof(hash32) );

            // Sign by OpenSSL
            unsigned char sigBuf[256]; // TEKITO
            unsigned int sigLen = 0;

            ret = RSA_sign( NID_sha256, hash32.bytes, 32, sigBuf, &sigLen, rsa_privkey );
            if (ret == 0)
            {
                printf( "error : RSA_sign(by OpenSSL)!\n" );
    			return 1;
            }
            printf("sigLen=%d\n", sigLen);
            PrintArray( (char*)"sigBuf", sigBuf, sigLen );


            // Verify for HSM
            ret = RSA_verify( NID_sha256, hash32.bytes, SHA256_DIGEST_LENGTH, middlePtr, middleLen, rsa_pubkey );
            if (ret == 0)
            {
                printf( "error : RSA_verify(by OpenSSL)!\n" );
    			return 1;
            }
            printf( "HSM Sign Verify Success!\n" );
        }
#endif

        PrintArray( (char*)"hash32", hash32.bytes, sizeof(hash32) );
//		PrintArray( (char*)"before", beforePtr, beforeLen );
		PrintArray( (char*)"middle", middlePtr, middleLen );

        return Status_OK;
    }

	// encrypt & dectypt test
	{
		unsigned char *beforePtr, *middlePtr, *afterPtr;
		int beforeLen, middleLen, afterLen;
		unsigned char *encPtr, *decPtr;
		struct NFast_Bignum *enc_bn, *dec_bn;
		M_Mech dec_mech;
				
		beforePtr = middlePtr = afterPtr = NULL;
		beforeLen = middleLen = afterLen = 0;
		encPtr = decPtr = NULL;
		enc_bn = dec_bn = NULL;
		
		// encrypt data setting
		beforeLen = DATA_LEN - 11;
		beforePtr = (unsigned char*)malloc( beforeLen );
		for ( i = 0; i < beforeLen; i++ )
			beforePtr[i] = ~i;
		my_bin2bignum( &(enc_bn), handle, beforePtr, beforeLen );

#if 0
		// my_bignum2bin test
		unsigned char debug_buf[ DATA_LEN ];
		PrintArray( (char*)"beforePtr", beforePtr, DATA_LEN );
		printf( "beforePtr -> bin2bignum -> bignum2bin -> debug_buf\n" );
		my_bignum2bin( debug_buf, DATA_LEN, handle, enc_bn );
		PrintArray( (char*)"debug_buf", debug_buf, DATA_LEN );
#endif 
		
		// encrypt transact
		cmd.cmd = Cmd_Encrypt;
		cmd.args.encrypt.flags = 0;	// Cmd_Encrypt_Args_flags_given_iv_present;
		cmd.args.encrypt.key = pub_keyid;
		cmd.args.encrypt.mech = Mech_RSApPKCS1;
		cmd.args.encrypt.plain.type = PlainTextType_Bignum;
		cmd.args.encrypt.plain.data.bignum.m = enc_bn;
		//cmd.args.encrypt.given_iv = NULL;
		result = NFastApp_Transact( nc, NULL, &cmd, &reply, NULL );
		if ( result != Status_OK )
		{	
			printf( "error(%d) : FastApp_Transact(Cmd_Encrypt)\n", result );
			return 1;
		}
		result = reply.status;
		if ( result != Status_OK )
		{	
			printf( "error(%d) : reply.status(Cmd_Encrypt)\n", result );
			return 1;
		}
#if 0
		if ( DATA_LEN != reply.reply.encrypt.cipher.data.rsappkcs1.m->nbytes )
		{
			printf( "error : output size isn't %d bytes!\n", DATA_LEN );
			return 1;
		}
#endif
		
		printf( "RSA data encrypt ok\n" );
		
		// decrypt data setting
		middleLen = reply.reply.encrypt.cipher.data.rsappkcs1.m->nbytes;
		middlePtr = (unsigned char*)malloc( middleLen );
		my_bignum2bin( middlePtr, middleLen, handle,
					reply.reply.encrypt.cipher.data.rsappkcs1.m );
		my_bin2bignum( &dec_bn, handle, middlePtr, middleLen );		
		dec_mech = reply.reply.encrypt.cipher.mech;
				
		NFastApp_Free_Command( handle, NULL, NULL, &cmd );
		NFastApp_Free_Reply( handle, NULL, NULL, &reply );
		memset( &cmd, 0, sizeof( cmd ) );
		memset( &reply, 0, sizeof( reply ) );

		// decrypt transact
		cmd.cmd = Cmd_Decrypt;
		cmd.args.decrypt.flags = 0;
		cmd.args.decrypt.key = priv_keyid;
		cmd.args.decrypt.mech = Mech_RSApPKCS1;
		cmd.args.decrypt.cipher.mech = dec_mech;
		cmd.args.decrypt.cipher.data.rsappkcs1.m = dec_bn;
		cmd.args.decrypt.reply_type = PlainTextType_Bignum;
		result = NFastApp_Transact( nc, NULL, &cmd, &reply, NULL );
		if ( result != Status_OK )
		{	
			printf( "error(%d) : FastApp_Transact(Cmd_Decrypt)\n", result );
			return 1;
		}
		result = reply.status;
		if ( result != Status_OK )
		{	
			printf( "error(%d) : reply.status(Cmd_Decrypt)\n", result );
			return 1;
		}
#if 0
		if ( DATA_LEN != reply.reply.decrypt.plain.data.bignum.m->nbytes )
		{
			printf( "error : output size isn't %d bytes!\n", DATA_LEN );
			return 1;
		}
#endif
		
		printf( "RSA data decrypt ok\n" );
		
		// after
		afterLen = reply.reply.encrypt.cipher.data.rsappkcs1.m->nbytes;
		afterPtr = (unsigned char*)malloc( afterLen );
		my_bignum2bin( afterPtr, afterLen, handle,
					reply.reply.decrypt.plain.data.bignum.m );
		
		NFastApp_Free_Command( handle, NULL, NULL, &cmd );
		NFastApp_Free_Reply( handle, NULL, NULL, &reply );
								
		// data show
		PrintArray( (char*)"before", beforePtr, beforeLen );
		PrintArray( (char*)"middle", middlePtr, middleLen );
		PrintArray( (char*)"after", afterPtr, afterLen );
	}	// encrypt & decrypt
		
	return result;
} // createRSAPSSCert

void PrintArray( char *pStr, const unsigned char *pData, int length )
{
	int i;
	printf( "%s(%d bytes)", pStr, length );
	for ( i = 0; i < length; i++ )
	{
		if ( (i % 16) == 0 ) printf( "\n" );
		printf( "%02X ", pData[ i ] );
	}
	printf( "\n" );
}	// PrintArray

int main( int argc, char *argv[] )
{
	int result = Status_OK;

    // initialize
    result = my_hsm_initialize( &handle, &nc, &pWorld, &card, &fips, &pModuleinfo );
	if ( result != Status_OK )
	{	
		printf( "error(%d) : my_hsm_initialize\n", result );
		return 1;
	}
		
    // create rsa-pss cert
    result = createRSAPSSCert( priv_keyident, pub_keyident );
	if ( result != Status_OK )
	{	
		printf( "error : createRSAPSSCert\n" );
		return 1;
	}
	printf( "RSA-PSS cert create success.\n" );

	// finalize
    result = my_hsm_finalize( &handle, &nc, &pWorld, &card, &fips );
	if ( result != Status_OK )
	{	
		printf( "error(%d) : my_hsm_finalize\n", result );
		return 1;
	}

	return 0;

} // main