#include "test_common.h"

static char *class_name[] = {
    "CKO_DATA",
    "CKO_CERTIFICATE",
    "CKO_PUBLIC_KEY",
    "CKO_PRIVATE_KEY",
    "CKO_SECRET_KEY",
    "CKO_HW_FEATURE",
    "CKO_DOMAIN_PARAMETERS",
    "CKO_VENDOR_DEFINED"
};

int
main(int argc, char *argv[])
{
    CK_RV             rc;
    CK_FLAGS          flags = 0;
    CK_SESSION_HANDLE sess;
    CK_ULONG          i;
    static CK_BBOOL   ltrue         = CK_TRUE;
    static CK_BBOOL   lfalse        = CK_FALSE;
    static const CK_OBJECT_CLASS oclass_pub   = CKO_PUBLIC_KEY;
    static const CK_OBJECT_CLASS oclass_priv  = CKO_PRIVATE_KEY;
    static CK_KEY_TYPE           key_type_256 = CKK_GOSTR3410;
    // PAR 256 A OID
    static CK_BYTE               gostR3410_256_params_A[] = {
        0x06, 0x07, 0x2a, 0x85, 0x03, 0x02, 0x02, 0x23, 0x01
    };
    // PAR 3411-2012-256 OID
    static CK_BYTE                      gostR3411_2012_256_params[] = {
        0x06, 0x08, 0x2a, 0x85, 0x03, 0x07, 0x01, 0x01, 0x02, 0x02
    };
    // Templates for token key pair generation.
    CK_ATTRIBUTE                       pub_2012_256_template[] = {
        { CKA_GOSTR3410PARAMS,
          gostR3410_256_params_A, sizeof(gostR3410_256_params_A) },
        { CKA_GOSTR3411PARAMS,
          gostR3411_2012_256_params, sizeof(gostR3411_2012_256_params) },
		{ CKA_TOKEN, &ltrue, sizeof(ltrue) }, 
    };
    CK_ATTRIBUTE                       priv_template[] = {
        { CKA_EXTRACTABLE, &lfalse, sizeof(lfalse) },
        { CKA_SENSITIVE,   &ltrue, sizeof(ltrue) },
		{ CKA_TOKEN, &ltrue, sizeof(ltrue) }, 
		{ CKA_PRIVATE, &ltrue, sizeof(ltrue) }, 
    };
    CK_MECHANISM                 mechanism_gen_256_desc =
    { CKM_GOSTR3410_KEY_PAIR_GEN, NULL, 0 };
    CK_MECHANISM_PTR             mechanism_gen_256      = &mechanism_gen_256_desc;
    CK_OBJECT_HANDLE             pub_key        = CK_INVALID_HANDLE;
    CK_OBJECT_HANDLE             priv_key       = CK_INVALID_HANDLE;
    CK_OBJECT_HANDLE             copy_priv_key   = CK_INVALID_HANDLE;
    // Template for hacking private key copy.
    CK_ATTRIBUTE                       copy_256_template[] = {
        { CKA_EXTRACTABLE, &ltrue, sizeof(ltrue) },
        { CKA_SENSITIVE,   &lfalse, sizeof(lfalse) },
		{ CKA_TOKEN, &lfalse, sizeof(lfalse) }, 
		{ CKA_PRIVATE, &lfalse, sizeof(lfalse) }, 
    };
	CK_BYTE priv_key_value[32];
	CK_ATTRIBUTE cka_value_template[] = {
		{ CKA_VALUE, NULL, 0 },
	};

    printf("CKA_SENSITIVE test\n");
    for (i = 1; i < (CK_ULONG) argc; i++)
    {
        if (strcmp("-api", argv[i]) == 0)
        {
            ++i;
            api_path = argv[i];
        }
        else if (strcmp("-user_pin", argv[i]) == 0)
        {
            ++i;
            user_pin = (CK_UTF8CHAR *)argv[i];
        }
    }
    rc = init(api_path, &funcs, NULL);
    if (rc != CKR_OK)
    {
        printf("init failed\n");
        return rc;
    }
    rc = get_slot_list(funcs,
                       &SlotList,
                       &SlotCount);
    if (rc != CKR_OK)
    {
        printf("get_slot_list failed, rc = 0x%lx\n", rc);
        return rc;
    }
    rc = find_slot_with_token(funcs,
                              SlotList,
                              SlotCount,
                              &SlotId);
    if (rc != CKR_OK)
    {
        printf("find_slot_with_token failed, rc = 0x%lx\n", rc);
        return rc;
    }

    flags = CKF_SERIAL_SESSION | CKF_RW_SESSION;
    rc    = funcs->C_OpenSession(SlotId, flags, NULL, NULL, &sess);
    if (rc != CKR_OK)
    {
        printf("C_OpenSession failed, rc = 0x%lx\n", rc);
        return rc;
    }

    rc = funcs->C_Login(sess, CKU_USER, user_pin, (CK_ULONG)strlen((char *)user_pin));
    if (rc != CKR_OK)
    {
        printf("C_Login failed, rc = 0x%lx\n", rc);
        return rc;
    }
    printf("C_Login (CKU_USER) success\n");

	printf("GOSTR3410-2012-256\n");
    rc = funcs->C_GenerateKeyPair(sess, mechanism_gen_256,
                                  pub_2012_256_template,
                                  sizeof(pub_2012_256_template) / sizeof(CK_ATTRIBUTE),
                                  priv_template, sizeof(priv_template) / sizeof(CK_ATTRIBUTE),
                                  &pub_key, &priv_key);
    if (rc != CKR_OK)
    {
        fprintf(stderr, "C_GenerateKeyPair failed, rc = 0x%lx\n", rc);
        goto end;
    }

	cka_value_template[0].pValue = priv_key_value;
	cka_value_template[0].ulValueLen = sizeof(priv_key_value);
	// Should fail for sensitive private key
	rc = funcs->C_GetAttributeValue(sess, priv_key, 
		cka_value_template, sizeof(cka_value_template)/sizeof(CK_ATTRIBUTE));
	fprintf(stderr, "C_GetAttributeValue: rc = 0x%lx\n", rc);
	if (rc == CKR_OK){
		printf("Base private key value:\n");
		print_bytes(cka_value_template[0].pValue, cka_value_template[0].ulValueLen);
	}
	// Should fail for sensitive private key
	rc = funcs->C_CopyObject(sess, priv_key, copy_256_template, 
		sizeof(copy_256_template)/sizeof(CK_ATTRIBUTE), &copy_priv_key); 
	fprintf(stderr, "C_CopyObject: rc = 0x%lx\n", rc);
	if (copy_priv_key != CK_INVALID_HANDLE){
		cka_value_template[0].pValue = priv_key_value;
		cka_value_template[0].ulValueLen = sizeof(priv_key_value);
		// Should fail for sensitive private key
		rc = funcs->C_GetAttributeValue(sess, copy_priv_key, 
			cka_value_template, sizeof(cka_value_template)/sizeof(CK_ATTRIBUTE));
		fprintf(stderr, "C_GetAttributeValue: rc = 0x%lx\n", rc);
		if (rc == CKR_OK){
			printf("Copied private key value:\n");
			print_bytes(cka_value_template[0].pValue, cka_value_template[0].ulValueLen);
		}
	}
end:
    printf("CKA_SENSITIVE test finished\n");
	if (copy_priv_key != CK_INVALID_HANDLE){
		funcs->C_DestroyObject(sess, copy_priv_key);
	}
	if (pub_key != CK_INVALID_HANDLE){
		funcs->C_DestroyObject(sess, pub_key);
	}
	if (priv_key != CK_INVALID_HANDLE){
		funcs->C_DestroyObject(sess, priv_key);
	}
    funcs->C_CloseSession(sess);

    return rc;
}
