#include <assert.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <getopt.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 DEFAULT_PRIV_KEY_FILE_PATH  "./test_key/test-rsa-privkey2048.der"
#define DEFAULT_PUB_KEY_FILE_PATH   "./test_key/test-rsa-pubkey2048.der"

#define SIGN_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" };

char *privKeyFilePath = (char*)DEFAULT_PRIV_KEY_FILE_PATH;
char *pubKeyFilePath = (char*)DEFAULT_PUB_KEY_FILE_PATH;
int isEnableOpenSSL = 0;
char *loadFilePath = NULL;
int loadFileSize = 0;
int validFileSize = 0;
int offsetFileSize = 0;
unsigned char *loadFileBuf = NULL;
char *writeFilePath = NULL;

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

int signVerifyRSAPKCS1Sha256( NFKM_KeyIdent priv_ident, NFKM_KeyIdent pub_ident )
{
	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

    // sign & verify
    {
        unsigned char *middlePtr = NULL;
		int middleLen = 0;
		struct NFast_Bignum *sign_bn = NULL;
		M_Mech verify_mech;
				
        // calc sha256
        M_Hash32 hash32;
        if (loadFileBuf)
        {
            SHA256( loadFileBuf + offsetFileSize, validFileSize, hash32.bytes );
        }
        else
        {
            int i;
            unsigned char buf[32];
            for (i = 0; i < sizeof(buf); i++)
            {
                buf[i] = ~i;
            }
            SHA256( buf, sizeof(buf), hash32.bytes );
        }

        // show data
        PrintArray( (char*)"Hash", hash32.bytes, sizeof(hash32.bytes) );

        // sign transact
        cmd.cmd = Cmd_Sign;
        cmd.args.sign.flags = 0;
        cmd.args.sign.key = priv_keyid;
        cmd.args.sign.mech = Mech_RSAhSHA256pPKCS1;
        cmd.args.sign.plain.type = PlainTextType_Hash32;
        cmd.args.sign.plain.data.hash32.data = hash32;

		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( "----- HSM 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( &sign_bn, handle, middlePtr, middleLen );		
        verify_mech = reply.reply.sign.sig.mech;

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

        // save sign data
        if (writeFilePath)
        {
            FILE *pSaveFile;

            pSaveFile = fopen( writeFilePath, "wb" );
            if (pSaveFile == NULL)
            {
                printf("file open error : %s\n", writeFilePath);
                return 1;
            }
            else
            {
                if (fwrite( middlePtr, middleLen, 1, pSaveFile ) != 1)
                {
                    printf("file write error : %s\n", writeFilePath);
                }
                else
                {
            		printf( "Success : Sign Data written to %s\n", writeFilePath );
                }
            }
            fclose( pSaveFile );
        }

		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_RSAhSHA256pPKCS1;
        cmd.args.verify.plain.type = PlainTextType_Hash32;
        cmd.args.verify.plain.data.hash32.data = hash32;
        cmd.args.verify.sig.mech = verify_mech;
        cmd.args.verify.sig.data.rsappkcs1.m = sign_bn;

		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( "----- HSM Verify OK -----\n" );

        // OpenSSL でも検証開始
        if (isEnableOpenSSL)
        {
            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( privKeyFilePath, &fileStat );
            privKeySize = fileStat.st_size;
            privKeyBuf = (unsigned char*)malloc( privKeySize );
            pPrivKeyFile = fopen( privKeyFilePath, "rb" );
            if (pPrivKeyFile == NULL)
            {
                printf("file open error\n");
                return 1;
            }
            fread( privKeyBuf, privKeySize, 1, pPrivKeyFile );
            fclose(pPrivKeyFile);

            // pub
            stat( pubKeyFilePath, &fileStat );
            pubKeySize = fileStat.st_size;
            pubKeyBuf = (unsigned char*)malloc( pubKeySize );
            pPubKeyFile = fopen( pubKeyFilePath, "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;

                // コマンドラインの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;
                }
            }

            // Sign by OpenSSL
            unsigned char signByOpenSSL[SIGN_LEN];
            unsigned int signLen = 0;

            ret = RSA_sign( NID_sha256, hash32.bytes, 32, signByOpenSSL, &signLen, rsa_privkey );
            if (ret == 0)
            {
                printf( "error : RSA_sign(by OpenSSL)!\n" );
    			return 1;
            }
            if (SIGN_LEN != signLen)
            {
                printf( "error : Sign Size(%d) by OpenSSL donen't match %d!\n", signLen, SIGN_LEN );
    			return 1;
            }
            PrintArray( (char*)"OpenSSL Sign ", signByOpenSSL, signLen );

            NFastApp_Free_Command( handle, NULL, NULL, &cmd );
    		NFastApp_Free_Reply( handle, NULL, NULL, &reply );

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

            // convert
            my_bin2bignum( &sign_bn, handle, signByOpenSSL, signLen );

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

		    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( "----- HSM Verify OK for OpenSSL Sign -----\n" );

            // 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( "----- OpenSSL Verify OK for HSM Sign -----\n" );

            free( privKeyBuf );
            free( pubKeyBuf );
        }

        return Status_OK;
    }
		
	return result;
} // signRSAPKCS1Sha256

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;

    int opt = 0;
    while (1)
    {
        int optionIndex = 0;
        static struct option longOptions[] =
        {
            { "openssl",    no_argument,        0,  0 },
            { "priv",       required_argument,  0,  0 },
            { "pub",        required_argument,  0,  0 },
            { "in",         required_argument,  0,  0 },
            { "out",        required_argument,  0,  0 },
            { "size",       required_argument,  0,  0 },
            { "offset",     required_argument,  0,  0 },
            { 0,            0,                  0,  0 }
        };
        // - によるオプション禁止、 -- のみ認める
        opt = getopt_long(argc, argv, "", longOptions, &optionIndex);
        if (opt == -1)
        {
            break;
        }
        
        switch (opt)
        {
        case 0 : // long arg
            {
                switch (optionIndex)
                {
                case 0 : // openssl
                    {
                        isEnableOpenSSL = 1;
                    } break;
                case 1 : // priv
                    {
                        privKeyFilePath = optarg;
                    } break;
                case 2 : // pub
                    {
                        pubKeyFilePath = optarg;
                    } break;
                case 3 : // in
                    {
                        loadFilePath = optarg;
                    } break;
                case 4 : // out
                    {
                        writeFilePath = optarg;
                    } break;
                case 5 : // size
                    {
                        validFileSize = atoi(optarg);
                    } break;
                case 6 : // offset
                    {
                        offsetFileSize = atoi(optarg);
                    } break;
                }
            } break;
        }
    }
    printf( "HSM sign & verify RSA PKCS1 SHA256 Sample.\n" );
    printf( " --openssl       : enable openssl sign & verify\n" );
    printf( " --priv filename : specify private RSA key (use openssl)\n" );
    printf( "  Default=%s\n", DEFAULT_PRIV_KEY_FILE_PATH );
    printf( " --pub filename  : specify public RSA key (use openssl)\n" );
    printf( "  Default=%s\n", DEFAULT_PUB_KEY_FILE_PATH );
    printf( " --in filename   : input file for SHA256\n" );
    printf( "  Default=reasonable array\n" );
    printf( " --out filename  : output sign data\n" );
    printf( "  Default=N/A\n" );
    printf( " --size size     : valid file size for SHA256\n" );
    printf( " --offset fize   : offset file size for SHA256\n" );
    printf( "\n" );
    printf( "----- start -----\n" );
    if (isEnableOpenSSL)
    {
        printf( "Enable OpenSSL Sign & Verify\n" );
        printf( "OpenSSL use privkey=%s\n", privKeyFilePath );
        printf( "OpenSSL use pubkey=%s\n", pubKeyFilePath );
    }
    if (loadFilePath)
    {
        FILE *pLoadFile;
        struct stat fileStat;

        stat( loadFilePath, &fileStat );
        loadFileSize = fileStat.st_size;
        loadFileBuf = (unsigned char*)malloc( loadFileSize );
        pLoadFile = fopen( loadFilePath, "rb" );
        if (pLoadFile == NULL)
        {
            printf("file open error : %s\n", loadFilePath);
            return 1;
        }
        fread( loadFileBuf, loadFileSize, 1, pLoadFile );
        fclose( pLoadFile );

        printf( "load filename=%s, size=%d bytes\n", loadFilePath, loadFileSize );
        if (validFileSize == 0 || validFileSize > loadFileSize)
        {
            validFileSize = loadFileSize;
        }
        if (offsetFileSize + validFileSize > loadFileSize)
        {
            offsetFileSize = loadFileSize - validFileSize;
        }
        printf( "valid size=%d bytes, offset=%d bytes\n", validFileSize, offsetFileSize );
    }

    // 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 pkcs1 sha256 sign & verify
    result = signVerifyRSAPKCS1Sha256( priv_keyident, pub_keyident );
	if ( result != Status_OK )
	{	
		printf( "error : createRSAPSSCert\n" );
		return 1;
	}
	printf( "RSA PKCS1 SHA256 Sign & Verify success.\n" );

    // free
    if (loadFileBuf)
    {
        free( loadFileBuf );
    }

	// 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