// Crypt.cpp : Defines the entry point for the console application.
//

#include <stdio.h>
#include <time.h>
#include <sys/stat.h>
#ifdef _WIN32
#include <windows.h>
#include <iostream>
#define CRYPT_VERIFYCONTEXT     0xF0000000
#define strlib "..\\lib\\strlib.dll"
#else
#define GetProcAddress dlsym
#define FreeLibrary dlclose
#include <string>
#include  <unistd.h>
#include  <dlfcn.h>
#define strlib "../lib/libstr.so.1.0"
#define _cdecl
#define CRYPT_VERIFYCONTEXT 0xF0000000
#define TRUE 1
#define FALSE 0
#endif
#ifndef _WIN32
typedef unsigned long DWORD;
typedef unsigned char BYTE;
typedef void *HINSTANCE;
typedef char CHAR;
#endif
typedef unsigned long HCRYPTPROV;
typedef unsigned long hCERTSTORE;
typedef unsigned long hCERTCOLLECTION;
typedef unsigned long hCRYPTMSG;
typedef unsigned long hCERTCONTEXT;
#define AT_KEYEXCHANGE          1
#define AT_SIGNATURE            2
// Константы для хранилища :
#define STORE_TYPE  "t=L;"
// Константы для шифрования:
#define DATA_OID    "1.2.840.113549.1.7.2"
#define CRYPT_OID   "1.3.6.1.4.1.6801.1.2.4"
#define DESCRIPTOR  "DESCRIPTOR"

typedef int(_cdecl *cspopencontext)
(HCRYPTPROV *phProv,
	CHAR *pszContainer,
	DWORD dwProvType,
	DWORD dwProvParam,
	DWORD dwFlags,
	BYTE * pbSignature,
	DWORD dwSigLen,
	BYTE * pbContainer,
	DWORD * dwContLen,
	HINSTANCE *phLib);
typedef int(_cdecl *crtopenstore)
(DWORD StoreProvider,
	HCRYPTPROV hProv,
	unsigned long dwFlags,
	bool  bFlag,
	const void *pvParam,
	hCERTSTORE *phStore);
typedef int(_cdecl *crtclosestore)    (hCERTSTORE *phStore);
typedef int(_cdecl *cspclosecontext)  (HCRYPTPROV *phProv);
typedef int(_cdecl *crtsetstorepropertyext)
(const hCERTSTORE hStore,
	const unsigned long Flags,
	const char *DName,
	const char *host,
	const char *port,
	const int timeOut,
	const char *name,
	const char *pass);
typedef char* (_cdecl *errorstring) (const int err);
typedef int(_cdecl *getmycertificate)
(hCERTSTORE hStore,
	const char *DName,
	const unsigned long keySpec,
	const bool isCA,
	unsigned char *certBody,
	unsigned long *bodySize);
typedef int(_cdecl *getcertificatecontext)
(const hCERTSTORE hStore,
	const unsigned char *certBody,
	const unsigned long bodySize,
	hCERTCONTEXT *phContext);
typedef int(_cdecl *crptmsgopen)
(HCRYPTPROV hProv,
	int Flags,
	hCRYPTMSG *phCrypt);
typedef int(_cdecl *getcertificatecollection)
(const hCERTSTORE hStore,
	const hCERTSTORE hMY,
	const char *keyOID,
	const int nameCount,
	char **DName,
	hCERTCOLLECTION *hCollection);
typedef int(_cdecl *cryptmessage)
(const hCRYPTMSG hCrypt,
	const hCERTCOLLECTION hCollection,
	const hCERTCONTEXT MYCert,
	const char *cryptOID,
	const char *dataOID,
	const bool Final,
	unsigned char *dataToEncrypted,
	unsigned long dataSize,
	const char *dataDescriptor,
	const unsigned long Flags,
	unsigned char *encryptedBlob,
	unsigned long *sizeBlob);
typedef int(_cdecl *bintobase)
(unsigned char *inBuf,
	long inSize,
	unsigned char *ouBuf,
	long *outSize);
typedef int(_cdecl *messagename)
(const unsigned char *cmsData,
	const unsigned long dataSize,
	int *senderCount,
	char **senderName,
	int *recipCount,
	char **recipName);
typedef int(_cdecl *releasecertificatecontext)
(const hCERTSTORE hStore,
	hCERTCONTEXT *phContext);
typedef int(_cdecl *crptmsgclose) (hCRYPTMSG *phCrypt);

cspopencontext            CSPOpenContext;
crtopenstore              CrtOpenStore;
crtclosestore             CrtCloseStore;
cspclosecontext           CSPCloseContext;
crtsetstorepropertyext    CrtSetStorePropertyExt;
errorstring               ErrorString;
getmycertificate          GetMYCertificate;
getcertificatecontext     GetCertificateContext;
crptmsgopen               CrptMsgOpen;
getcertificatecollection  GetCertificateCollection;
cryptmessage              CryptMessage;
bintobase                 BinToBase;
messagename               MessageName;
releasecertificatecontext ReleaseCertificateContext;
crptmsgclose              CrptMsgClose;

void WriteFile(unsigned char * data);
void ReadFile(char* file, unsigned char **data, long *size);
HINSTANCE load_lib(char *lib)
{
#ifdef _WIN32
	return LoadLibraryA(lib);
#else
	return dlopen(lib, RTLD_LAZY);
#endif
}
void lib_free(HINSTANCE Instance)
{
#ifdef _WIN32
	FreeLibrary(Instance);
#else
	dlclose(Instance);
#endif
}
int get_func(HINSTANCE lib)
{
	CSPOpenContext = (cspopencontext)GetProcAddress(lib, "CSPOpenContext");
	if (!CSPOpenContext)
	{
		printf("error load function - %s", "CSPOpenContext");
		return 1;
	}
	CrtOpenStore = (crtopenstore)GetProcAddress(lib, "CrtOpenStore");
	if (!CrtOpenStore)
	{
		printf("error load function - %s", "CrtOpenStore");
		return 1;
	}
	CSPCloseContext = (cspclosecontext)GetProcAddress(lib, "CSPCloseContext");
	if (!CSPCloseContext)
	{
		printf("error load function - %s", "CSPCloseContext");
		return 1;
	}
	CrtCloseStore = (crtclosestore)GetProcAddress(lib, "CrtCloseStore");
	if (!CrtCloseStore)
	{
		printf("error load function - %s", "CrtCloseStore");
		return 1;
	}
	ErrorString = (errorstring)GetProcAddress(lib, "ErrorString");
	if (!ErrorString)
	{
		printf("error load function - %s", "ErrorString");
		return 1;
	}
	GetMYCertificate = (getmycertificate)GetProcAddress(lib, "GetMYCertificate");
	if (!GetMYCertificate)
	{
		printf("error load function - %s", "GetMYCertificate");
		return 1;
	}
	GetCertificateContext = (getcertificatecontext)GetProcAddress(lib, "GetCertificateContext");
	if (!GetCertificateContext)
	{
		printf("error load function - %s", "GetCertificateContext");
		return 1;
	}
	CrptMsgOpen = (crptmsgopen)GetProcAddress(lib, "CrptMsgOpen");
	if (!CrptMsgOpen)
	{
		printf("error load function - %s", "CrptMsgOpen");
		return 1;
	}
	GetCertificateCollection = (getcertificatecollection)GetProcAddress(lib, "GetCertificateCollection");
	if (!GetCertificateCollection)
	{
		printf("error load function - %s", "GetCertificateCollection");
		return 1;
	}
	CryptMessage = (cryptmessage)GetProcAddress(lib, "CryptMessage");
	if (!CryptMessage)
	{
		printf("error load function - %s", "CryptMessage");
		return 1;
	}
	BinToBase = (bintobase)GetProcAddress(lib, "BinToBase");
	if (!BinToBase)
	{
		printf("error load function - %s", "BinToBase");
		return 1;
	}
	MessageName = (messagename)GetProcAddress(lib, "MessageName");
	if (!MessageName)
	{
		printf("error load function - %s", "MessageName");
		return 1;
	}
	CrtSetStorePropertyExt = (crtsetstorepropertyext)GetProcAddress(lib, "CrtSetStorePropertyExt");
	if (!CrtSetStorePropertyExt)
	{
		printf("error load function - %s", "CrtSetStorePropertyExt");
		return 1;
	}
	ReleaseCertificateContext = (releasecertificatecontext)GetProcAddress(lib, "ReleaseCertificateContext");
	if (!ReleaseCertificateContext)
	{
		printf("error load function - %s", "ReleaseCertificateContext");
		return 1;
	}
	CrptMsgClose = (crptmsgclose)GetProcAddress(lib, "CrptMsgClose");
	if (!CrptMsgClose)
	{
		printf("error load function - %s", "CrptMsgClose");
		return 1;
	}
	return 0;
}
int main(int argc, char* argv[])
{
	HINSTANCE lib = NULL;
	int ret = 0;
	int code = 0;
	int i = 0;
	char host[128];
	char port[32];
	char DN[256];
	char My_DN[256];
	char datafile[256];
	char profile[128];
	HCRYPTPROV hP = 0;
	hCERTSTORE hStore = 0;
	unsigned long certSize = 0;
	unsigned char *myBlob = NULL;
	hCERTCONTEXT hContext = 0;
	hCRYPTMSG hCrypt = 0;
	char **DNames = NULL;
	hCERTCOLLECTION hCollection = 0;
	unsigned char *encryptBlob = NULL;
	unsigned char *base64 = NULL;
	unsigned long sizeBlob = 0;
	long baseBlob = 0;
	int  senderCount = 0;
	int  recipCount = 0;
	char **senderName = NULL;
	char **recipName = NULL;
	unsigned char *data = NULL;
	long dataSize = 0;
	
	setlocale(LC_ALL, "Russian");
	
	if (argc != 7)
	{
		printf("%s store_host store_port my_profile my_dn recip file\n", argv[0]);
		return ret;
	}
	for (i = 1; i<argc; i++)
	{
		if (i == 1) sprintf(host, "%s\0", argv[i]);
		if (i == 2) sprintf(port, "%s\0", argv[i]);
		if (i == 3)
#ifdef _WIN32
			sprintf(profile, "profile://%s\0", argv[i]);
#else
			sprintf(profile, "profile://%s", argv[i]);
#endif
		if (i == 4) sprintf(My_DN, "%s\0", argv[i]);
		if (i == 5) sprintf(DN, "%s\0", argv[i]);
		if (i == 6) sprintf(datafile, "%s\0", argv[i]);
	}
	ReadFile(datafile, &data, &dataSize);
	if (!data) throw 1;
	printf("read data file - ok\n");
	try
	{
		lib = load_lib(strlib);
		if (get_func(lib))
		{
			printf("error load func.\n");
			exit(1);
		}
		//получить контекст криптопровайдера
		code = CSPOpenContext(&hP, profile, 25, 0, 0, NULL, 0, NULL, 0, NULL);
		if (code)
		{
			printf("%s [CSPOpenContext: error = %d]\n", ErrorString(code), code);
			throw code;
		}
		printf("CSPOpenContext - ok\n");
		//открыть хранилище сертификатов
		code = CrtOpenStore(0, hP, 0, false, NULL, &hStore);
		if (code)
		{
			printf("%s [CrtOpenStore: error = %d]\n", ErrorString(code), code);
			throw code;
		}
		printf("CrtOpenStore - ok\n");
		//настроить хранилище сертификатов
		code = CrtSetStorePropertyExt(hStore, 0, STORE_TYPE, host, port, 300, 0, 0);
		if (code)
		{
			printf("%s [CrtSetStorePropertyExt: error = %i\n]", ErrorString(code), code);
			throw code;
		}
		printf("CrtSetStorePropertyExt - ok\n");
		//получить собственный сертификат пользователя
		//перед вызовом функции обязательно вызывать CrtSetStorePropertyExt, т.к. GetMYCertificate обращается к хранилищю LDAP, но
		//если собственный сертификат будет загружен из файла, то можно сразу вызывать GetCertificateContext
		code = GetMYCertificate(hStore, My_DN, AT_KEYEXCHANGE, false, NULL, &certSize);
		if (code)
		{
			printf("%s [GetMYCertificate: error = %d]\n", ErrorString(code), code);
			throw code;
		}
		myBlob = (unsigned char*)calloc(certSize, sizeof(unsigned char));
		code = GetMYCertificate(hStore, My_DN, AT_KEYEXCHANGE, false, myBlob, &certSize);
		if (code)
		{
			printf("%s [GetMYCertificate: error = %d]\n", ErrorString(code), code);
			throw code;
		}
		printf("GetMYCertificate - ok\n");
		//получить контекст сертификата
		code = GetCertificateContext(hStore, myBlob, certSize, &hContext);
		if (code)
		{
			printf("%s [GetCertificateContext: error = %d]\n", ErrorString(code), code);
			throw code;
		}
		printf("GetCertificateContext - ok\n");
		//установить хендл на класс CMS-crypt
		code = CrptMsgOpen(hP, 0, &hCrypt);
		if (code)
		{
			printf("%s [CrptMsgOpen: error = %d]\n", ErrorString(code), code);
			throw code;
		}
		printf("CrptMsgOpen - ok\n");
		//получить хендл на коллекцию сертификатов
		DNames = (char**)calloc(1, sizeof(char*));
		DNames[0] = (char*)calloc(2048, sizeof(char));
		sprintf(DNames[0], DN);
		code = GetCertificateCollection(hStore, hStore, "1.2.398.3.10.1.4.1.2", 1, DNames, &hCollection);
		if (code)
		{
			printf("%s [GetCertificateCollection: error = %d]\n", ErrorString(code), code);
			throw code;
		}
		printf("GetCertificateCollection - ok\n");
		//формирование зашифрованного сообщения
		//получить размер зашифрованного сообщения
		code = CryptMessage(hCrypt, hCollection, hContext, (char*)CRYPT_OID, (char*)DATA_OID, true, data, dataSize, (char*)DESCRIPTOR, 1, NULL, &sizeBlob);
		if (code)
		{
			printf("%s [CryptMessage: error = %d]\n", ErrorString(code), code);
			throw code;
		}
		encryptBlob = (unsigned char*)calloc(sizeBlob, sizeof(unsigned char));
		//сформировать зашифрованное сообщение
		code = CryptMessage(hCrypt, hCollection, hContext, (char*)CRYPT_OID, (char*)DATA_OID, true, data, dataSize, (char*)DESCRIPTOR, 1, encryptBlob, &sizeBlob);
		if (code)
		{
			printf("%s [CryptMessage: error = %d]\n", ErrorString(code), code);
			throw code;
		}
		base64 = (unsigned char*)calloc(sizeBlob * 2, sizeof(unsigned char));
		code = BinToBase(encryptBlob, (long)sizeBlob, base64, &baseBlob);
		if (code)
		{
			printf("%s [BintoBase: error = %d]\n", ErrorString(code), code);
			throw code;
		}
		WriteFile(base64);
		printf("CryptMessage - ok\n");
		//получить информацию об отправителях/получателях зашифрованного сообщения
		code = MessageName(encryptBlob, sizeBlob, &senderCount, NULL, &recipCount, NULL);
		if (code)
		{
			printf("%s [MessageName: error = %d]\n", ErrorString(code), code);
			throw code;
		}
		if (senderCount)
		{
			senderName = (char**)calloc(senderCount, sizeof(char*));
			for (i = 0; i<senderCount; i++)
				senderName[i] = (char*)calloc(260, sizeof(char));
		}
		if (recipCount)
		{
			recipName = (char**)calloc(recipCount, sizeof(char*));
			for (i = 0; i<recipCount; i++)
				recipName[i] = (char*)calloc(260, sizeof(char));
		}
		code = MessageName(encryptBlob, sizeBlob, &senderCount, senderName, &recipCount, recipName);
		if (code)
		{
			printf("%s [MessageName: error = %d]\n", ErrorString(code), code);
			throw code;
		}
		for (i = 0; i<senderCount; i++)
			printf("sender %d name - %s\n", i + 1, senderName[i]);
		for (i = 0; i<recipCount; i++)
			printf("recip  %d name - %s\n", i + 1, recipName[i]);
		printf("Programme succeeds, code = 0\n");
	}
	catch (int er) { ret = er; }
	if (myBlob) { free(myBlob); myBlob = NULL; }
	if (DNames) { if (DNames[0]) free(DNames[0]); free(DNames); }
	if (encryptBlob) { free(encryptBlob); encryptBlob = NULL; }
	if (base64) { free(base64); base64 = NULL; }
	if (data) { delete[] data; data = NULL; }
	if (senderName)
	{
		for (i = 0; i<senderCount; i++)
		{
			free(senderName[i]); senderName[i] = NULL;
		}
		free(senderName); senderName = NULL;
	}
	if (recipName)
	{
		for (i = 0; i<recipCount; i++)
		{
			free(recipName[i]); recipName[i] = NULL;
		}
		free(recipName); recipName = NULL;
	}
	//очистить kонтекст сертификата
	if (hContext)
	{
		ReleaseCertificateContext(hStore, &hContext);
		hContext = (hCERTCONTEXT)NULL;
	}
	//закрыть хранилище сертификатов
	if (hStore)
	{
		CrtCloseStore(&hStore);
		hStore = (hCERTSTORE)NULL;
	}
	//удалить контекст криптопровайдера
	if (hP)
	{
		CSPCloseContext(&hP);
		hP = (HCRYPTPROV)NULL;
	}
	//закрыть хендл на класс CMS-crypt
	if (hCrypt)
	{
		CrptMsgClose(&hCrypt);
		hCrypt = (hCRYPTMSG)NULL;
	}
	if (lib) lib_free(lib);
	return ret;
}
//------------------------------------------------------------------------------
void WriteFile(unsigned char * data)
{
	FILE *fl = NULL;
	fl = fopen("cms", "w+");
	if (!fl)
	{
		return;
		printf("error - can not open file\n");
	}
	fprintf(fl, "%s", data);
	fclose(fl);
}
//------------------------------------------------------------------------------
void ReadFile(char* file, unsigned char **data, long *size)
{
	FILE *fl = NULL;
	struct stat  st;
	if (stat(file, &st)) return;
	*size = st.st_size;
	*data = new unsigned char[st.st_size + 1];
	memset(*data, 0, *size + 1);
	if ((fl = fopen(file, "rb")) == NULL)
	{
		delete *data;
		*size = 0;
		*data = NULL;
		return;
	}
	fread(*data, *size, 1, fl);
	fclose(fl);
	return;
}
