//   a utility to import PFX files into the local machine store
//   by Keith Brown
//   http://www.develop.com/kbrown
//
//   a free utility for you to use at your own risk
//
#define UNICODE
#define _UNICODE
#include <windows.h>
#include <wincrypt.h>
#include <stdio.h>
#include <string>
#include <iostream>

#pragma comment(lib, "crypt32.lib")

void _err(const wchar_t* fcn, DWORD err = GetLastError()) 
{
	wchar_t szErrMsg[256];
	wchar_t szMsg[512];
	if ( FormatMessage( FORMAT_MESSAGE_FROM_SYSTEM, 0, err, 0, szErrMsg, sizeof szErrMsg / sizeof *szErrMsg, 0 ) )
	{
		 swprintf( szMsg, L"%s failed: %s", fcn, szErrMsg );
	}else 
	{
		swprintf( szMsg, L"%s failed: 0x%08X", fcn, err );
	}
	wprintf(L"%s\n", szMsg);
}

void getPassword(wchar_t* pwd, int cch) 
{
	pwd[0] = L'\0';
	if (cch < 128)
		return;

	// turn off console echo
    HANDLE hConsole = GetStdHandle(STD_INPUT_HANDLE);
    DWORD oldMode;
    if (!GetConsoleMode(hConsole, &oldMode)) 
    {
        _err(L"GetConsoleMode");
		return;
    }
    UINT newMode = oldMode & ~ENABLE_ECHO_INPUT;
    if (!SetConsoleMode(hConsole, newMode)) 
    {
        _err(L"SetConsoleMode");
		return;
    }

	wprintf(L"Enter password:\n");
	if (1 != wscanf(L"%128ls", pwd)) 
	{
		pwd[0] = L'\0';		
	}
	wprintf(L"\n");

    // restore console echo
    SetConsoleMode(hConsole, oldMode);
}

void printUsage()
{
	std::cout << "Usage: certimport [-e] [-p] password filename\n";
}

int wmain(int argc, wchar_t* argv[]) 
{
	HANDLE hfile = INVALID_HANDLE_VALUE;
	HANDLE hsection = 0;
	void* pfx = 0;
	HCERTSTORE pfxStore = 0;
	HCERTSTORE myStore = 0;
	HCERTSTORE rootStore = 0;
	bool exportable = false;

	wchar_t password[128];
	wcscpy(password, L"foobar");
	wchar_t filename[128];

	// Handle all arguments.
	// We accept the following arguments:
	// -p password: If given, use the password, else use "foobar".
	// -e: If given use CRYPT_EXPORTABLE.
	// last string MUST be the filename. 
	 
	int optionIndex = 1;
	while ((optionIndex< argc) && (argv[optionIndex][0]=='-')) 
	{
		std::wstring sw = argv[optionIndex];
		if (sw==L"-e") 
		{
			exportable = true;
		}
		else if (sw==L"-p") 
		{
			optionIndex++;
			wcscpy(password, argv[optionIndex]);
		}
		else
		{
			std::cout << "Unknown switch: " << argv[optionIndex] << std::endl;
			printUsage();
			return 1;
		}
		optionIndex++;
	}
		
	if(optionIndex < argc)
	{
		wcscpy(filename , argv[optionIndex]);
	}else
	{
		printUsage();
		return 1;
	}
	
	hfile = CreateFile(	filename, 
						FILE_READ_DATA, 
						FILE_SHARE_READ, 
						0, 
						OPEN_EXISTING, 
						0, 
						0);
	if (INVALID_HANDLE_VALUE == hfile) 
	{
		_err(L"CreateFile");
		goto cleanup;
	}

	hsection = CreateFileMapping(	hfile, 
									0, 
									PAGE_READONLY, 
									0, 
									0, 
									0);
	if (!hsection) 
	{
		_err(L"CreateFileMapping");
		goto cleanup;
	}

	pfx = MapViewOfFile(	hsection, 
							FILE_MAP_READ, 
							0, 
							0, 
							0);
	if (!pfx) 
	{
		_err(L"MapViewOfFile");
		goto cleanup;
	}

	CRYPT_DATA_BLOB blob;
	blob.cbData = GetFileSize(	hfile, 
								0);
	blob.pbData = (BYTE*)pfx;
	
	if (!PFXIsPFXBlob(&blob)) 
	{
		wprintf(L"%s is not a valid PFX file\n", filename);
		goto cleanup;
	}

	//getPassword(password, sizeof password / sizeof *password);
	
	DWORD importFlags = CRYPT_MACHINE_KEYSET;
	if (exportable) 
	{
		importFlags |= CRYPT_EXPORTABLE;
	}
	
	pfxStore = PFXImportCertStore(	&blob, 
									password, 
									importFlags);
	if (!pfxStore) 
	{
		_err(L"PFXImportCertStore");
		goto cleanup;
	}

	myStore = CertOpenStore(	CERT_STORE_PROV_SYSTEM, 
								0, 
								0, 
								CERT_STORE_OPEN_EXISTING_FLAG | CERT_SYSTEM_STORE_LOCAL_MACHINE, 
								L"MY");
	if (!myStore) 
	{
		_err(L"CertOpenSystemStore MY");
		goto cleanup;
	}
	rootStore = CertOpenStore(	CERT_STORE_PROV_SYSTEM, 
								0, 
								0, 
								CERT_STORE_OPEN_EXISTING_FLAG | CERT_SYSTEM_STORE_LOCAL_MACHINE, 
								L"Root");
	if (!rootStore ) 
	{
		_err(L"CertOpenSystemStore Root");
		goto cleanup;
	}
	
	// We assume two certificates, the first one being the root certificate,
	// all other certificates personal.
	unsigned long counter = 0;
	
	PCCERT_CONTEXT pctx = 0;
	while (0 != (pctx = CertEnumCertificatesInStore(pfxStore, pctx))) 
	{
		wchar_t name[128];
		if (CertGetNameString(pctx, CERT_NAME_FRIENDLY_DISPLAY_TYPE, 0, 0, name, sizeof name / sizeof *name)) 
		{
			wprintf(L"Found a certificate in the PFX file: %s\n", name);
		}else 
		{
			_err(L"CertGetNameString");
		}
		wprintf(L"Attempting to import certificate into machine store...\n");

		if (CertAddCertificateContextToStore(counter == 0 ? rootStore : myStore, pctx, CERT_STORE_ADD_NEW, 0)) 
		{
			wprintf(L"Import succeeded.\n");
		}else 
		{
			DWORD err = GetLastError();
			if (CRYPT_E_EXISTS == err) 
			{
				wprintf(L"\nAn equivalent certificate already exists. Overwrite? (y/n) ");
				wchar_t yesOrNo;
				scanf("%lc", &yesOrNo);
				scanf("%lc", &yesOrNo);
				if (L'y' == yesOrNo) 
				{
					if (CertAddCertificateContextToStore(counter == 0 ? rootStore : myStore, pctx, CERT_STORE_ADD_REPLACE_EXISTING, 0)) 
					{
						wprintf(L"Import succeeded.\n");
					}else 
					{
						_err(L"CertAddCertificateContextToStore");
						goto cleanup;
					}
				}else 
				{
					wprintf(L"Skipped.\n");
				}
			}else 
			{
				_err(L"CertAddCertificateContextToStore");
				goto cleanup;
			}
		}
		++counter;
	}

cleanup:
	if (myStore) 
	{
		CertCloseStore(myStore, 0);
	}
	if (pfxStore) 
	{
		CertCloseStore(pfxStore, CERT_CLOSE_STORE_FORCE_FLAG);
	}
	if (pfx) 
	{
		UnmapViewOfFile(pfx);
	}
	if (hsection) 
	{
		CloseHandle(hsection);
	}
	if (INVALID_HANDLE_VALUE != hfile) 
	{
		CloseHandle(hfile);
	}
	return 0;
}