#include "stdafx.h"

#include "Crypt.hpp"

#include <assert.h>

#include <vector>



CCryptException::CCryptException(DWORD errorcode)
	: errorcode_(errorcode)
{
}

DWORD CCryptException::GetErrorCode() const
{
	return errorcode_;
}


//////////////////////////////////////////

CCrypt::CCrypt()
	: hProv_(0)
{
}

CCrypt::~CCrypt() throw()
{
	// プロバイダを解除する。
	if (hProv_) {
		::CryptReleaseContext(hProv_, 0);
		hProv_ = NULL;
	}
}

/*!
 * プロバイダを初期化します。
 * すでに初期化済みであれば何もしません。
 */
void CCrypt::init()
{
	if (hProv_) {
		// すでに初期化済みであれば何もしない。
		return;
	}

	// 暗号化プロバイダの取得
	// Windows98/200のためにAESではなくDESとする。
	// Win98のためにCRYPT_VERIFYCONTEXTで秘密キー等のキーストアを使用しないことを明示する。
	if ( !CryptAcquireContext(&hProv_, NULL, NULL, PROV_RSA_FULL, CRYPT_VERIFYCONTEXT)) {
		throw CCryptException(GetLastError());
	}

	// サポートされているアルゴリズムの一覧の取得
	loadSupportedAlgorithm();
}

/*!
 * サポートされているアルゴリズムの一覧を取得します。
 */
void CCrypt::loadSupportedAlgorithm()
{
	assert(hProv_ && "プロバイダが初期化されていません。");

	algtypemap_.clear();

	PROV_ENUMALGS alginfo = {0};

	DWORD enumflg = CRYPT_FIRST;
	for (;;) {
		DWORD alginfosiz = sizeof(PROV_ENUMALGS);
		BOOL ret = CryptGetProvParam(
			hProv_,
			PP_ENUMALGS,
			(BYTE*) &alginfo,
			&alginfosiz,
			enumflg);

		if (!ret) {
			DWORD errorcode = GetLastError();
			if (errorcode == ERROR_NO_MORE_ITEMS) {
				break;
			}
			throw CCryptException(errorcode);
		}

		// アルゴリズムIDと名前のペアを格納
		CA2T algName(alginfo.szName);
		algtypemap_.insert(algtype(alginfo.aiAlgid, tstring(algName)));

		enumflg = CRYPT_NEXT;
	}
}


const CCrypt::algtype& CCrypt::chooseAlgorithm()
{
	// 初期化
	init();
	assert(hProv_ && "プロバイダが初期化されていません。");

	// 優先順
	ALG_ID algidPriority[] = {CALG_3DES_112, CALG_3DES, CALG_DES, 0};
	
	ALG_ID *pAlg = algidPriority;
	while (*pAlg) {
		algtypemap::const_iterator ite = algtypemap_.find(*pAlg);
		if (ite != algtypemap_.end()) {
			return *ite;
		}
		pAlg++;
	}
	throw CCryptException(ERROR_INVALID_PARAMETER);
}

/*!
 * 指定した名前の暗号化アルゴリズムをサポートしている中から取得する。
 * 該当がない場合は例外が発生する。
 * \param アルゴリズム名
 * \return アルゴリズム
 */
const CCrypt::algtype& CCrypt::findAlgorithm(const tstring& algName)
{
	// 初期化
	init();
	assert(hProv_ && "プロバイダが初期化されていません。");

	for (algtypemap::const_iterator ite = algtypemap_.begin();
		ite != algtypemap_.end();
		++ite) {
		// 名前と一致する暗号化アルゴリズムを探し出す
		if (ite->second == algName) {
			return *ite;
		}
	}
	throw CCryptException(ERROR_INVALID_PARAMETER);
}

void CCrypt::encrypt(const tstring& password, const algtype &algtyp, const tstring& src, tstring& result)
{
	if (password.empty()) {
		// パスワードが空の場合は、そのまま返す。
		result = src;
		return;
	}

	// 初期化
	init();
	assert(hProv_ && "プロバイダが初期化されていません。");

	// パスワードベースの対称キー生成
	CCryptPassword passwordKey(*this);
	if ( !passwordKey.setPassword(algtyp, password)) {
		// 暗号化キーが生成されなかった場合は平文のまま返す
		result = src;
		return;
	}
	assert(passwordKey.hPasswordKey_ && "暗号化キーが設定されていません。");

	// 暗号化バッファにソースを転記する。
	// (ソースが空であっても暗号化する。)
	LPCTSTR szSrc = src.c_str();
	DWORD bytelen = static_cast<DWORD>(sizeof(TCHAR) * src.length());
	DWORD bufsiz = ((bytelen / 16) + 1) * 16; // 128bit block size
	
	std::vector<unsigned char> buf(bufsiz);
	memcpy(&buf[0], szSrc, bytelen);

	// 暗号化
	if ( !CryptEncrypt(passwordKey.hPasswordKey_, NULL, TRUE, 0, &buf[0], &bytelen, bufsiz)) {
		throw CCryptException(GetLastError());
	}

	// 暗号結果の先頭にアルゴリズム名を設定
	result.clear();
	result += _T("{{");
	result += algtyp.second;
	result += _T("}}");

	// 暗号結果を16進数で表現
	TCHAR szHex[10] = {0};
	for (int idx = 0; idx < (int) bytelen; idx++) {
		unsigned int ch = buf[idx];
		_stprintf_s(szHex, 10, _T("%02x"), ch);
		result += szHex;
	}
}

void CCrypt::decrypt(const tstring& password, const tstring& src, tstring& result)
{
	if (password.empty()) {
		// パスワードが空の場合は、そのまま返す。
		result = src;
		return;
	}
	if (src.empty()) {
		// ソースから空行であれば、そのまま返す。
		result.clear();
		return;
	}

	tstring head = _T("{{");
	size_t endpt = src.npos;
	if (src.compare(0, head.length(), head) != 0
		|| (endpt = src.find(_T("}}"))) == src.npos) {
		// {{で始まっていない場合、もしくは}}がない場合は
		// 暗号化名がないので平文とみなす。
		result = src;
		return;
	}

	// 初期化
	init();

	// アルゴリズム名とHEX部の分離
	tstring algName = src.substr(2, endpt - 2);
	tstring strHex = src.substr(endpt + 2);

	algtype algtyp = findAlgorithm(algName);

	// パスワードベースの対称キーの生成
	CCryptPassword passwordKey(*this);
	if ( !passwordKey.setPassword(algtyp, password)) {
		// 暗号化キーが生成されなかったらソースを、そのまま返す。
		result = src;
		return;
	}
	assert(passwordKey.hPasswordKey_ && "暗号化キーが設定されていません。");

	// hexからバイト列に変換
	TCHAR szHex[3] = {0};

	std::vector<unsigned char> buf;

	LPCTSTR p = strHex.c_str();
	while (*p) {
		szHex[0] = *p++;
		szHex[1] = *p ? *p++ : 0;
		int byte = _tcstol(szHex, NULL, 16);
		buf.push_back((unsigned char) byte);
	}
	DWORD bytelen = static_cast<DWORD>(buf.size());

	// 暗号化解除
	if ( !CryptDecrypt(passwordKey.hPasswordKey_, NULL, TRUE, 0, &buf[0], &bytelen)) {
		throw CCryptException(GetLastError());
	}

	// 終端処理
	buf.push_back(0);
	buf[bytelen] = 0;
#ifdef _UNICODE
	buf.push_back(0);
	buf[bytelen + 1] = 0;
#endif

	// バイト列から文字列へ
	result.clear();
	result = tstring((LPCTSTR) &buf[0]);
}



///////////////////////////////////////////////////////


CCrypt::CCryptPassword::CCryptPassword(CCrypt &crypt)
	: crypt_(crypt)
	, hPasswordHash_(0)
	, hPasswordKey_(0)
{
}

CCrypt::CCryptPassword::~CCryptPassword() throw()
{
	// パスワードクリアする。
	clearPassword();
}

/*!
 * 初期化し、プロバイダのハンドルを取得します。
 * すでに初期化済みであれば、そのハンドルを返します。
 */
HCRYPTPROV CCrypt::CCryptPassword::init()
{
	crypt_.init();
	assert(crypt_.hProv_ && "暗号化プロバイダが初期化されていません。");
	return crypt_.hProv_;
}

/*!
 * パスワードを設定します。
 * 空の場合はパスワードが解除されます。
 */
bool CCrypt::CCryptPassword::setPassword(const CCrypt::algtype &algtyp, const tstring& password)
{
	// 初期化されていなければ初期化する。
	HCRYPTPROV hProv = init();

	// すでにパスワードが設定済みであればクリアする。
	clearPassword();

	if (password.empty()) {
		// パスワードが空であれば、ハッシュとキーは不要。
		return false;
	}

	HCRYPTHASH hHash = NULL;
	HCRYPTKEY hKey= NULL;
	try {
		// ハッシュを構築する。
		hHash = createHash(hProv, password);

		// 対称暗号化キーを生成する。
		hKey = createKey(hProv, hHash, algtyp);
	}
	catch (...) {
		if (hKey) {
			::CryptDestroyKey(hKey);
		}
		if (hHash) {
			::CryptDestroyHash(hHash);
		}
		throw;
	}

	assert(hHash && "ハッシュハンドルが設定されていません。");
	assert(hKey && "暗号化キーが設定されていません。");

	hPasswordHash_ = hHash;
	hPasswordKey_ = hKey;
	algName_ = algtyp.second;

	return true;
}

/*!
 * パスワードを解除します。
 */
void CCrypt::CCryptPassword::clearPassword()
{
	if (hPasswordKey_) {
		::CryptDestroyKey(hPasswordKey_);
		hPasswordKey_ = NULL;
	}

	if (hPasswordHash_) {
		::CryptDestroyHash(hPasswordHash_);
		hPasswordHash_ = NULL;
	}
}

/*!
 * パスワードベースのハッシュを生成して返します。
 * \param password パスワード
 * \return ハッシュ
 */
HCRYPTHASH CCrypt::CCryptPassword::createHash(HCRYPTPROV hProv, const tstring& password)
{
	assert(hProv && "プロバイダが初期化されていません。");

	HCRYPTHASH hHash = 0;
	if ( !CryptCreateHash(hProv, CALG_SHA, 0, 0, &hHash)) {
		throw CCryptException(GetLastError());
	}
	try {
		LPCTSTR szPassword = password.c_str();
		DWORD bytelen = static_cast<DWORD>(sizeof(TCHAR) * password.length());
		if ( !CryptHashData(hHash, (const BYTE *) szPassword, bytelen, 0)) {
			throw CCryptException(GetLastError());
		}
	}
	catch (...) {
		::CryptDestroyHash(hHash);
		hHash = 0;
		throw;
	}

	return hHash;
}

/*!
 * ハッシュから対称暗号化キーを生成します。
 * \param hHash ハッシュ
 * \param alg アルゴリズム
 * \return 暗号化キー
 */
HCRYPTKEY CCrypt::CCryptPassword::createKey(HCRYPTPROV hProv, HCRYPTHASH hHash, const CCrypt::algtype &alg)
{
	assert(hProv && "プロバイダが初期化されていません。");
	assert(hHash && "ハッシュがNULLです。");

	ALG_ID algid = alg.first;
	HCRYPTKEY hCryptKey = 0;
	if ( !CryptDeriveKey(hProv, algid, hHash, CRYPT_NO_SALT, &hCryptKey)) {
		throw CCryptException(GetLastError());
	}
	return hCryptKey;
}

