#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include "../../include/cryptoki_ext.h"
#include "exportcert.h"
#include <iostream>
using namespace std;

CExportcert::CExportcert()
{
	m_pSlotList = NULL_PTR;
	m_hSession = NULL_PTR;
}

CExportcert::~CExportcert()
{
	if(m_hSession)
	{
		C_CloseSession(m_hSession);
		m_hSession = NULL_PTR;
	}
	if(m_pSlotList)
	{
		delete[] m_pSlotList;
		m_pSlotList = NULL_PTR;
	}
}

CK_RV CExportcert::Init() 
{
	if(m_hSession) return CKR_OK;
	
	CK_RV rv;
	CK_ULONG ulCount = 0;
	rv = C_GetSlotList(TRUE, NULL_PTR, &ulCount);
	if(CKR_OK != rv )
	{
		printf("Can't aquire Token Information\n");
		return rv;
	}
	if(0 >= ulCount)
	{
		printf("Can't connect to Token, Make sure you have inserted Token.\n");
		return CKR_GENERAL_ERROR;
	}
	
	m_pSlotList = (CK_SLOT_ID_PTR)new CK_SLOT_ID[ulCount];
	if (! m_pSlotList) 
	{
		printf("can't allocate enough memory\n");
		return CKR_HOST_MEMORY;
	}
	
	rv = C_GetSlotList(TRUE, m_pSlotList, &ulCount);
	if(CKR_OK != rv )
	{
		printf("Can't aquire Token Information\n");
		return rv;
	}
	if(0 >= ulCount)
	{
		printf("Can't connect to Token, Make sure you have inserted Token.\n");
		return CKR_GENERAL_ERROR;
	}
	
	rv = C_OpenSession(
		m_pSlotList[0],  CKF_RW_SESSION | CKF_SERIAL_SESSION,
		NULL_PTR, NULL_PTR, &m_hSession);
	if(CKR_OK != rv )
	{
		printf("Can't aquire Token Information\n");
		delete[] m_pSlotList;
		m_pSlotList = NULL_PTR;
		return rv;
	}
	return rv;
}

void CExportcert::Export() 
{
	CK_OBJECT_CLASS dataClass = CKO_CERTIFICATE;
	bool IsToken=true;
	CK_ATTRIBUTE pTempl[] = 
	{
		{CKA_CLASS, &dataClass, sizeof(CKO_CERTIFICATE)},
		{CKA_TOKEN, &IsToken, sizeof(true)}
	};
	
	C_FindObjectsInit(m_hSession, pTempl, 2);

	CK_OBJECT_HANDLE hCKObj;
	CK_ULONG ulRetCount = 0;
	CK_RV ckrv = 0;
	CK_ULONG conut = 0;
	CK_LONG i = 0;
	do
	{
		ckrv = C_FindObjects(m_hSession, &hCKObj, 1, &ulRetCount);
		if(CKR_OK != ckrv)
		{
			break;	
		}
		if(1 != ulRetCount)
		{
			break;
		}
		
		CK_ATTRIBUTE pAttrTemp[] = 
		{
			{CKA_CLASS, NULL, 0},
			{CKA_CERTIFICATE_TYPE,NULL,0},
			{CKA_LABEL, NULL, 0},
			{CKA_SUBJECT,NULL,0},
			{CKA_ID,NULL,0},
			{CKA_VALUE,NULL,0}
		};
		
		ckrv = C_GetAttributeValue(m_hSession, hCKObj, pAttrTemp, 6);
		if(ckrv != CKR_OK)
		{
			break;
		}
		
		pAttrTemp[0].pValue = new char[pAttrTemp[0].ulValueLen];
		pAttrTemp[1].pValue = new char[pAttrTemp[1].ulValueLen];
		pAttrTemp[2].pValue = new char[pAttrTemp[2].ulValueLen+1];
		pAttrTemp[3].pValue = new char[pAttrTemp[3].ulValueLen+1];
		pAttrTemp[4].pValue = new char[pAttrTemp[4].ulValueLen+1];
		pAttrTemp[5].pValue = new char[pAttrTemp[5].ulValueLen ];
	
		memset(pAttrTemp[0].pValue, 0, pAttrTemp[0].ulValueLen);
		memset(pAttrTemp[1].pValue, 0, pAttrTemp[1].ulValueLen);
		memset(pAttrTemp[2].pValue, 0, pAttrTemp[2].ulValueLen+1);	
		memset(pAttrTemp[3].pValue, 0, pAttrTemp[3].ulValueLen+1);
		memset(pAttrTemp[4].pValue, 0, pAttrTemp[4].ulValueLen+1);
		memset(pAttrTemp[5].pValue, 0, pAttrTemp[5].ulValueLen);
		
		ckrv = C_GetAttributeValue(m_hSession, hCKObj, pAttrTemp, 6);
		if(ckrv != CKR_OK)
		{
			delete[](char*) pAttrTemp[0].pValue;
			delete[](char*) pAttrTemp[1].pValue;
			delete[](char*) pAttrTemp[2].pValue;
			delete[](char*) pAttrTemp[3].pValue;
			delete[](char*) pAttrTemp[4].pValue;
			delete[](char*) pAttrTemp[5].pValue;
			break;
		}
		
		char *strvale = (char*)pAttrTemp[2].pValue;
		int fd;
		char s[256];
		i++;
		int j,h = strlen(strvale);
		for(j = 0; j < h;j++)
		{
			if(strvale[j] == ' ')
				strvale[j] = '_';
		}
		sprintf(s,"%s%d.cer",strvale,i);
		fd = open(s,O_RDWR|O_CREAT,0644);
		if(0 > fd)
		{
			printf("Open file fault\n");
			return ;
		}
		write(fd, pAttrTemp[5].pValue, pAttrTemp[5].ulValueLen);
		close(fd);
		printf("Export successful.\n");

		delete[](char*) pAttrTemp[0].pValue;
		delete[](char*) pAttrTemp[1].pValue;
		delete[](char*) pAttrTemp[2].pValue;
		delete[](char*) pAttrTemp[3].pValue;
		delete[](char*) pAttrTemp[4].pValue;
		delete[](char*) pAttrTemp[5].pValue;
		conut++;
	}while(true);
	if(conut > 1)
	{
		printf("Export %d certificates\n", conut);
	}else
	{
		printf("Export %d certificate\n", conut);
	}
}
