#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <ctype.h>
#include <time.h>

#include "rng.h"
#include "api.h"
#include "kaz_api.h"
#include "gmp.h"

#define	MAX_MARKER_LEN		50

#define KAT_SUCCESS          0
#define KAT_FILE_OPEN_ERROR -1
#define KAT_DATA_ERROR      -3
#define KAT_CRYPTO_FAILURE  -4

int		FindMarker(FILE *infile, const char *marker);
int		ReadHex(FILE *infile, unsigned char *A, int Length, char *str);
void	fprintBstr(FILE *fp, char *S, unsigned char *A, unsigned long long L);

char    AlgName[] = "KAZ-KA-256";

int main(void)
{
    char                   fn_req[32], fn_rsp[32];
    FILE                   *fp_req, *fp_rsp;
    unsigned char          *shka, *shkb;
    int                    count;
    int                    done;
    unsigned char          pka[KAZ_KA_PUBLICKEY_BYTES*2], ska[KAZ_KA_PRIVATEKEY_BYTES*2];
    unsigned char          pkb[KAZ_KA_PUBLICKEY_BYTES*2], skb[KAZ_KA_PRIVATEKEY_BYTES*2];
    int                    ret_val;

    clock_t tkeygenA, tkeygenAtotal=0, tsharedA, tsharedAtotal=0;
    clock_t tkeygenB, tkeygenBtotal=0, tsharedB, tsharedBtotal=0;

    // Create the REQUEST file
    sprintf(fn_req, "PQCkaKAT_%d.req", KAZ_KA_PRIVATEKEY_BYTES*2);
    if ( (fp_req = fopen(fn_req, "w")) == NULL ) {
        printf("Couldn't open <%s> for write\n", fn_req);
        return KAT_FILE_OPEN_ERROR;
    }
    sprintf(fn_rsp, "PQCkaKAT_%d.rsp", KAZ_KA_PRIVATEKEY_BYTES*2);
    if ( (fp_rsp = fopen(fn_rsp, "w")) == NULL ) {
        printf("Couldn't open <%s> for write\n", fn_rsp);
        return KAT_FILE_OPEN_ERROR;
    }
	
    for (int i=0; i<100; i++) {
        fprintf(fp_req, "count = %d\n", i);
        fprintf(fp_req, "pka =\n");
        fprintf(fp_req, "ska =\n");
        fprintf(fp_req, "pkb =\n");
        fprintf(fp_req, "skb =\n");
        fprintf(fp_req, "shka =\n\n");
        fprintf(fp_req, "shkb =\n\n");
    }
    fclose(fp_req);

    //Create the RESPONSE file based on what's in the REQUEST file
    if ( (fp_req = fopen(fn_req, "r")) == NULL ) {
        printf("Couldn't open <%s> for read\n", fn_req);
        return KAT_FILE_OPEN_ERROR;
    }

    fprintf(fp_rsp, "# %s\n\n", CRYPTO_ALGNAME);
    done = 0;

    do {
        if ( FindMarker(fp_req, "count = ") )
            fscanf(fp_req, "%d", &count);
        else {
            done = 1;
            break;
        }
        fprintf(fp_rsp, "count = %d\n", count);

        shka = (unsigned char *)calloc(KAZ_KA_SHAREDKEY_BYTES, sizeof(unsigned char));
        shkb = (unsigned char *)calloc(KAZ_KA_SHAREDKEY_BYTES, sizeof(unsigned char));

        tkeygenA=clock();
        // Generate the public/private keypair A
        if ( (ret_val = crypto_ka_keypair(pka, ska, "PRIMARY")) != 0) {
            printf("crypto_ka_keypair returned <%d>\n", ret_val);
            return KAT_CRYPTO_FAILURE;
        }
        tkeygenA=clock()-tkeygenA;
        tkeygenAtotal+=tkeygenA;
        tkeygenB=clock();
//printf("PASS.\n");
        // Generate the public/private keypair B
        if ( (ret_val = crypto_ka_keypair(pkb, skb, NULL)) != 0) {
            printf("crypto_ka_keypair returned <%d>\n", ret_val);
            return KAT_CRYPTO_FAILURE;
        }
//printf("PASS.\n");
        tkeygenB=clock()-tkeygenB;
        tkeygenBtotal+=tkeygenB;
        fprintBstr(fp_rsp, "pka = ", pka, KAZ_KA_PUBLICKEY_BYTES*2);
        fprintBstr(fp_rsp, "ska = ", ska, KAZ_KA_PRIVATEKEY_BYTES*2);
        fprintBstr(fp_rsp, "pkb = ", pkb, KAZ_KA_PUBLICKEY_BYTES*2);
        fprintBstr(fp_rsp, "skb = ", skb, KAZ_KA_PRIVATEKEY_BYTES*2);

        tsharedA=clock();
        // Generate the Shared Key AB
        if ( (ret_val = crypto_ka(shka, pkb, ska)) != 0 ) {
            printf("crypto_ka returned <%d>\n", ret_val);
            return KAT_CRYPTO_FAILURE;
        }
        tsharedA=clock()-tsharedA;
        tsharedAtotal+=tsharedA;
        tsharedB=clock();
//printf("PASS.\n");
        // Generate the Shared Key BA
        if ( (ret_val = crypto_ka(shkb, pka, skb)) != 0 ) {
            printf("crypto_ka returned <%d>\n", ret_val);
            return KAT_CRYPTO_FAILURE;
        }
//printf("PASS.\n");
        tsharedB=clock()-tsharedB;
        tsharedBtotal+=tsharedB;
        fprintBstr(fp_rsp, "shka = ", shka, KAZ_KA_SHAREDKEY_BYTES);
        fprintBstr(fp_rsp, "shkb = ", shkb, KAZ_KA_SHAREDKEY_BYTES);
        fprintf(fp_rsp, "\n");

        if ( memcmp(shka, shkb, KAZ_KA_SHAREDKEY_BYTES) ) {
            printf("crypto_key_agreement returned FALSE.\n");
            return KAT_CRYPTO_FAILURE;
        }

        memset(pka, 0, KAZ_KA_PUBLICKEY_BYTES*2);
        memset(ska, 0, KAZ_KA_PRIVATEKEY_BYTES*2);
        memset(pkb, 0, KAZ_KA_PUBLICKEY_BYTES*2);
        memset(skb, 0, KAZ_KA_PRIVATEKEY_BYTES*2);
        free(shka);
        free(shkb);

    } while ( !done );

    printf("Approx milliseconds, keygen A: %.0f\n", (double)((double)tkeygenAtotal / CLOCKS_PER_SEC) * 1000);
    printf("Approx milliseconds, keygen B: %.0f\n", (double)((double)tkeygenBtotal / CLOCKS_PER_SEC) * 1000);
    printf("Approx milliseconds, sharedkey A: %.0f\n", (double)((double)tsharedAtotal / CLOCKS_PER_SEC) * 1000);
    printf("Approx milliseconds, sharedkey B: %.0f\n", (double)((double)tsharedBtotal / CLOCKS_PER_SEC) * 1000);

    fclose(fp_req);
    fclose(fp_rsp);

    return KAT_SUCCESS;
}

//
// ALLOW TO READ HEXADECIMAL ENTRY (KEYS, DATA, TEXT, etc.)
//
int
FindMarker(FILE *infile, const char *marker)
{
	char	line[MAX_MARKER_LEN];
	int		i, len;
	int curr_line;

	len = (int)strlen(marker);
	if ( len > MAX_MARKER_LEN-1 )
		len = MAX_MARKER_LEN-1;

	for ( i=0; i<len; i++ )
	  {
	    curr_line = fgetc(infile);
	    line[i] = curr_line;
	    if (curr_line == EOF )
	      return 0;
	  }
	line[len] = '\0';

	while ( 1 ) {
		if ( !strncmp(line, marker, len) )
			return 1;

		for ( i=0; i<len-1; i++ )
			line[i] = line[i+1];
		curr_line = fgetc(infile);
		line[len-1] = curr_line;
		if (curr_line == EOF )
		    return 0;
		line[len] = '\0';
	}

	// shouldn't get here
	return 0;
}

//
// ALLOW TO READ HEXADECIMAL ENTRY (KEYS, DATA, TEXT, etc.)
//
int
ReadHex(FILE *infile, unsigned char *A, int Length, char *str)
{
	int			i, ch, started;
	unsigned char	ich;

	if ( Length == 0 ) {
		A[0] = 0x00;
		return 1;
	}
	memset(A, 0x00, Length);
	started = 0;
	if ( FindMarker(infile, str) )
		while ( (ch = fgetc(infile)) != EOF ) {
			if ( !isxdigit(ch) ) {
				if ( !started ) {
					if ( ch == '\n' )
						break;
					else
						continue;
				}
				else
					break;
			}
			started = 1;
			if ( (ch >= '0') && (ch <= '9') )
				ich = ch - '0';
			else if ( (ch >= 'A') && (ch <= 'F') )
				ich = ch - 'A' + 10;
			else if ( (ch >= 'a') && (ch <= 'f') )
				ich = ch - 'a' + 10;
            else // shouldn't ever get here
                ich = 0;

			for ( i=0; i<Length-1; i++ )
				A[i] = (A[i] << 4) | (A[i+1] >> 4);
			A[Length-1] = (A[Length-1] << 4) | ich;
		}
	else
		return 0;

	return 1;
}

void
fprintBstr(FILE *fp, char *S, unsigned char *A, unsigned long long L)
{
	unsigned long long  i;

	fprintf(fp, "%s", S);

	for ( i=0; i<L; i++ )
		fprintf(fp, "%02X", A[i]);

	if ( L == 0 )
		fprintf(fp, "00");

	fprintf(fp, "\n");
}
