pwmanager.cpp

来自「一款密码保险箱源码」· C++ 代码 · 共 2,201 行 · 第 1/5 页

	ASSERT(pEntry->pszPassword[pEntry->uPasswordLen] == 0);
}

void CPwManager::UnlockEntryPassword(__inout_ecount(1) PW_ENTRY *pEntry)
{
	ASSERT_ENTRY(pEntry);
	LockEntryPassword(pEntry); // OFB encryption mode

#ifdef _DEBUG
	unsigned int i;
	for(i = 0; i < pEntry->uPasswordLen; i++)
	{
		ASSERT(pEntry->pszPassword[i] != 0);
	}
#endif
}

void CPwManager::NewDatabase()
{
	_DeleteEntryList(TRUE); // Delete really everything, the strings too
	_DeleteGroupList(TRUE);

	m_pLastEditedEntry = NULL;

	_AllocGroups(PWM_NUM_INITIAL_GROUPS); // Allocate some space for the new items
	_AllocEntries(PWM_NUM_INITIAL_ENTRIES);

	m_vUnknownMetaStreams.clear();
}

#define _OPENDB_FAIL_LIGHT \
{ \
	if(pVirtualFile != NULL) \
	{ \
		mem_erase((unsigned char *)pVirtualFile, uAllocated); \
		SAFE_DELETE_ARRAY(pVirtualFile); \
	} \
	m_dwKeyEncRounds = PWM_STD_KEYENCROUNDS; \
}
#define _OPENDB_FAIL \
{ \
	_OPENDB_FAIL_LIGHT; \
	SAFE_DELETE_ARRAY(pwGroupTemplate.pszGroupName); \
	SAFE_DELETE_ARRAY(pwEntryTemplate.pszTitle); \
	SAFE_DELETE_ARRAY(pwEntryTemplate.pszURL); \
	SAFE_DELETE_ARRAY(pwEntryTemplate.pszUserName); \
	SAFE_DELETE_ARRAY(pwEntryTemplate.pszPassword); \
	SAFE_DELETE_ARRAY(pwEntryTemplate.pszAdditional); \
	SAFE_DELETE_ARRAY(pwEntryTemplate.pszBinaryDesc); \
	SAFE_DELETE_ARRAY(pwEntryTemplate.pBinaryData); \
	return PWE_INVALID_FILESTRUCTURE; \
}

#define RESET_TIME_FIELD_NORMAL(pTimeEx) { \
	(pTimeEx)->btDay = 1; (pTimeEx)->btHour = 0; (pTimeEx)->btMinute = 0; \
	(pTimeEx)->btMonth = 1; (pTimeEx)->btSecond = 0; (pTimeEx)->shYear = 2004; }
#define RESET_TIME_FIELD_EXPIRE(pTimeEx) { \
	(pTimeEx)->btDay = 28; (pTimeEx)->btHour = 23; (pTimeEx)->btMinute = 59; \
	(pTimeEx)->btMonth = 12; (pTimeEx)->btSecond = 59; (pTimeEx)->shYear = 4092; }

#define RESET_PWG_TEMPLATE(ptrx) { \
	memset(ptrx, 0, sizeof(PW_GROUP)); \
	RESET_TIME_FIELD_NORMAL(&(ptrx)->tCreation); RESET_TIME_FIELD_NORMAL(&(ptrx)->tLastMod); \
	RESET_TIME_FIELD_NORMAL(&(ptrx)->tLastAccess); RESET_TIME_FIELD_EXPIRE(&(ptrx)->tExpire); }
#define RESET_PWE_TEMPLATE(ptrx) { \
	memset(ptrx, 0, sizeof(PW_ENTRY)); \
	RESET_TIME_FIELD_NORMAL(&(ptrx)->tCreation); RESET_TIME_FIELD_NORMAL(&(ptrx)->tLastMod); \
	RESET_TIME_FIELD_NORMAL(&(ptrx)->tLastAccess); RESET_TIME_FIELD_EXPIRE(&(ptrx)->tExpire); }

// int CPwManager::OpenDatabase(const TCHAR *pszFile, __out_opt PWDB_REPAIR_INFO *pRepair)
// {
//	return this->OpenDatabaseEx(pszFile, pRepair, NULL);
// }

// If bIgnoreCorrupted is TRUE the manager will try to ignore all database file
// errors, i.e. try to read as much as possible instead of breaking out at the
// first error.
// To open a file normally, set bIgnoreCorrupted to FALSE (default).
// To open a file in rescue mode, set it to TRUE.
int CPwManager::OpenDatabase(const TCHAR *pszFile, __out_opt PWDB_REPAIR_INFO *pRepair)
{
	FILE *fp;
	char *pVirtualFile;
	unsigned long uFileSize, uAllocated, uEncryptedPartSize;
	unsigned long pos;
	PW_DBHEADER hdr;
	sha256_ctx sha32;
	UINT8 uFinalKey[32];
	char *p;
	char *pStart;
	USHORT usFieldType;
	DWORD dwFieldSize;
	PW_GROUP pwGroupTemplate;
	PW_ENTRY pwEntryTemplate;

	ASSERT(sizeof(char) == 1);

	ASSERT(pszFile != NULL); if(pszFile == NULL) return PWE_INVALID_PARAM;
	ASSERT(pszFile[0] != 0); if(pszFile[0] == 0) return PWE_INVALID_PARAM; // Length != 0

	RESET_PWG_TEMPLATE(&pwGroupTemplate);
	RESET_PWE_TEMPLATE(&pwEntryTemplate);

	if(pRepair != NULL) { ZeroMemory(pRepair, sizeof(PWDB_REPAIR_INFO)); }

	fp = NULL;
	_tfopen_s(&fp, pszFile, _T("rb"));
	if(fp == NULL) return PWE_NOFILEACCESS_READ;

	// Get file size
	fseek(fp, 0, SEEK_END);
	uFileSize = ftell(fp);
	fseek(fp, 0, SEEK_SET);

	if(uFileSize < sizeof(PW_DBHEADER))
		{ fclose(fp); return PWE_INVALID_FILEHEADER; }

	// Allocate enough memory to hold the complete file
	uAllocated = uFileSize + 16 + 1 + 8 + 4; // 16 = encryption buffer space, 1+8 = string terminating NULL (UTF-8), 4 unused
	pVirtualFile = new char[uAllocated];
	if(pVirtualFile == NULL) { fclose(fp); return PWE_NO_MEM; }
	memset(&pVirtualFile[uFileSize + 17 - 1], 0, 1 + 8);

	fread(pVirtualFile, 1, uFileSize, fp);
	fclose(fp);

	// Extract header structure from memory file
	memcpy(&hdr, pVirtualFile, sizeof(PW_DBHEADER));

	// Check if we can open this
	if((hdr.dwSignature1 != PWM_DBSIG_1) || (hdr.dwSignature2 != PWM_DBSIG_2))
		{ _OPENDB_FAIL_LIGHT; return PWE_INVALID_FILESIGNATURE; }

	if((hdr.dwVersion & 0xFFFFFF00) != (PWM_DBVER_DW & 0xFFFFFF00))
	{
		if((hdr.dwVersion == 0x00020000) || (hdr.dwVersion == 0x00020001) || (hdr.dwVersion == 0x00020002))
		{
			if(pVirtualFile != NULL)
			{
				mem_erase((unsigned char *)pVirtualFile, uAllocated);
				SAFE_DELETE_ARRAY(pVirtualFile);
			}
			return (_OpenDatabaseV2(pszFile) != FALSE) ? PWE_SUCCESS : PWE_UNKNOWN;
		}
		else if(hdr.dwVersion <= 0x00010002)
		{
			if(pVirtualFile != NULL)
			{
				mem_erase((unsigned char *)pVirtualFile, uAllocated);
				SAFE_DELETE_ARRAY(pVirtualFile);
			}
			return (_OpenDatabaseV1(pszFile) != FALSE) ? PWE_SUCCESS : PWE_UNKNOWN;
		}
		else { ASSERT(FALSE); _OPENDB_FAIL; }
	}

	// Select algorithm
	if((hdr.dwFlags & PWM_FLAG_RIJNDAEL) != 0) m_nAlgorithm = ALGO_AES;
	else if((hdr.dwFlags & PWM_FLAG_TWOFISH) != 0) m_nAlgorithm = ALGO_TWOFISH;
	else { ASSERT(FALSE); _OPENDB_FAIL; }

	m_dwKeyEncRounds = hdr.dwKeyEncRounds;

	// Generate m_pTransformedMasterKey from m_pMasterKey
	if(_TransformMasterKey(hdr.aMasterSeed2) == FALSE) { ASSERT(FALSE); _OPENDB_FAIL; }

	// Hash the master password with the salt in the file
	sha256_begin(&sha32);
	sha256_hash(hdr.aMasterSeed, 16, &sha32);
	sha256_hash(m_pTransformedMasterKey, 32, &sha32);
	sha256_end((unsigned char *)uFinalKey, &sha32);

	if(pRepair == NULL)
	{
		// ASSERT(((uFileSize - sizeof(PW_DBHEADER)) % 16) == 0);
		if(((uFileSize - sizeof(PW_DBHEADER)) % 16) != 0)
		{
			_OPENDB_FAIL_LIGHT;
			return PWE_INVALID_FILESIZE;
		}
	}
	else // Repair the database
	{
		if(((uFileSize - sizeof(PW_DBHEADER)) % 16) != 0)
		{
			uFileSize -= sizeof(PW_DBHEADER); ASSERT((uFileSize & 0xF) != 0);
			uFileSize &= ~0xF;
			uFileSize += sizeof(PW_DBHEADER);
		}

		ASSERT(((uFileSize - sizeof(PW_DBHEADER)) % 16) == 0);

		pRepair->dwOriginalGroupCount = hdr.dwGroups;
		pRepair->dwOriginalEntryCount = hdr.dwEntries;
	}

	if(m_nAlgorithm == ALGO_AES)
	{
		CRijndael aes;

		// Initialize Rijndael algorithm
		if(aes.Init(CRijndael::CBC, CRijndael::DecryptDir, uFinalKey,
			CRijndael::Key32Bytes, hdr.aEncryptionIV) != RIJNDAEL_SUCCESS)
			{ _OPENDB_FAIL_LIGHT; return PWE_CRYPT_ERROR; }

		// Decrypt! The first bytes aren't encrypted (that's the header)
		uEncryptedPartSize = (unsigned long)aes.PadDecrypt((UINT8 *)pVirtualFile + sizeof(PW_DBHEADER),
			uFileSize - sizeof(PW_DBHEADER), (UINT8 *)pVirtualFile + sizeof(PW_DBHEADER));
	}
	else if(m_nAlgorithm == ALGO_TWOFISH)
	{
		CTwofish twofish;

		if(twofish.Init(uFinalKey, 32, hdr.aEncryptionIV) != true)
			{ _OPENDB_FAIL };

		uEncryptedPartSize = (unsigned long)twofish.PadDecrypt((UINT8 *)pVirtualFile + sizeof(PW_DBHEADER),
			uFileSize - sizeof(PW_DBHEADER), (UINT8 *)pVirtualFile + sizeof(PW_DBHEADER));
	}
	else
	{
		ASSERT(FALSE); _OPENDB_FAIL; // This should never happen
	}

#if 0
	// For debugging purposes, a file containing the plain text is created.
	// This code of course must not be compiled into the final binary.
#pragma message("PLAIN TEXT OUTPUT IS ENABLED!")
#pragma message("DO NOT DISTRIBUTE THIS BINARY!")
	// std::basic_string<TCHAR> tstrClear = pszFile;
	// tstrClear += _T(".plaintext.bin");
	// FILE *fpClear = NULL;
	// _tfopen_s(&fpClear, tstrClear.c_str(), _T("wb"));
	// fwrite(pVirtualFile, 1, uFileSize, fpClear);
	// fclose(fpClear); fpClear = NULL;
#endif

	// Check for success (non-repair mode only)
	if(pRepair == NULL)
	{
		if((uEncryptedPartSize > 2147483446) || ((uEncryptedPartSize == 0) &&
			((hdr.dwGroups != 0) || (hdr.dwEntries != 0))))
		{
			_OPENDB_FAIL_LIGHT;
			return PWE_INVALID_KEY;
		}
	}

	// Check if key is correct (with very high probability)
	if(pRepair == NULL)
	{
		sha256_begin(&sha32);
		sha256_hash((unsigned char *)pVirtualFile + sizeof(PW_DBHEADER), uEncryptedPartSize, &sha32);
		sha256_end((unsigned char *)uFinalKey, &sha32);
		if(memcmp(hdr.aContentsHash, uFinalKey, 32) != 0)
			{ _OPENDB_FAIL_LIGHT; return PWE_INVALID_KEY; }
	}

	NewDatabase(); // Create a new database and initialize internal structures

	// Add groups from the memory file to the internal structures
	unsigned long uCurGroup;
	BOOL bRet;
	pos = sizeof(PW_DBHEADER);
	pStart = &pVirtualFile[pos];
	for(uCurGroup = 0; uCurGroup < hdr.dwGroups; )
	{
		p = &pVirtualFile[pos];

		if(pRepair != NULL) if(IsBadReadPtr(p, 2) != FALSE) { _OPENDB_FAIL; }
		memcpy(&usFieldType, p, 2);
		p += 2; pos += 2;
		if(pos >= uFileSize) { _OPENDB_FAIL; }

		if(pRepair != NULL) if(IsBadReadPtr(p, 4) != FALSE) { _OPENDB_FAIL; }
		memcpy(&dwFieldSize, p, 4);
		p += 4; pos += 4;
		if(pos >= (uFileSize + dwFieldSize)) { _OPENDB_FAIL; }

		if(pRepair != NULL) if(IsBadReadPtr(p, dwFieldSize) != FALSE) { _OPENDB_FAIL; }
		bRet = ReadGroupField(usFieldType, dwFieldSize, (BYTE *)p, &pwGroupTemplate);
		if((usFieldType == 0xFFFF) && (bRet == TRUE))
			uCurGroup++; // Now and ONLY now the counter gets increased

		p += dwFieldSize;
		if(p < pStart) { _OPENDB_FAIL; }
		pos += dwFieldSize;
		if(pos >= uFileSize) { _OPENDB_FAIL; }
	}
	SAFE_DELETE_ARRAY(pwGroupTemplate.pszGroupName);

	// Get the entries
	unsigned long uCurEntry;
	for(uCurEntry = 0; uCurEntry < hdr.dwEntries; )
	{
		p = &pVirtualFile[pos];

		if(pRepair != NULL) if(IsBadReadPtr(p, 2) != FALSE) { _OPENDB_FAIL; }
		memcpy(&usFieldType, p, 2);
		p += 2; pos += 2;
		if(pos >= uFileSize) { _OPENDB_FAIL; }

		if(pRepair != NULL) if(IsBadReadPtr(p, 4) != FALSE) { _OPENDB_FAIL; }
		memcpy(&dwFieldSize, p, 4);
		p += 4; pos += 4;
		if(pos >= (uFileSize + dwFieldSize)) { _OPENDB_FAIL; }

		if(pRepair != NULL) if(IsBadReadPtr(p, dwFieldSize) != FALSE) { _OPENDB_FAIL; }
		bRet = ReadEntryField(usFieldType, dwFieldSize, (BYTE *)p, &pwEntryTemplate);
		if((usFieldType == 0xFFFF) && (bRet == TRUE))
			uCurEntry++; // Now and ONLY now the counter gets increased

		p += dwFieldSize;
		if(p < pStart) { _OPENDB_FAIL; }
		pos += dwFieldSize;
		if(pos >= uFileSize) { _OPENDB_FAIL; }
	}
	SAFE_DELETE_ARRAY(pwEntryTemplate.pszTitle);
	SAFE_DELETE_ARRAY(pwEntryTemplate.pszURL);
	SAFE_DELETE_ARRAY(pwEntryTemplate.pszUserName);
	SAFE_DELETE_ARRAY(pwEntryTemplate.pszPassword);
	SAFE_DELETE_ARRAY(pwEntryTemplate.pszAdditional);
	SAFE_DELETE_ARRAY(pwEntryTemplate.pszBinaryDesc);
	SAFE_DELETE_ARRAY(pwEntryTemplate.pBinaryData);

	memcpy(&m_dbLastHeader, &hdr, sizeof(PW_DBHEADER));

	// Erase and delete memory file
	mem_erase((unsigned char *)pVirtualFile, uAllocated);
	SAFE_DELETE_ARRAY(pVirtualFile);

	DWORD dwRemovedStreams = _LoadAndRemoveAllMetaStreams(true);
	if(pRepair != NULL) pRepair->dwRecognizedMetaStreamCount = dwRemovedStreams;
	VERIFY(DeleteLostEntries() == 0);
	FixGroupTree();

	return PWE_SUCCESS;
}

int CPwManager::SaveDatabase(const TCHAR *pszFile)
{
	FILE *fp;
	char *pVirtualFile;
	DWORD uFileSize, uEncryptedPartSize, uAllocated;
	DWORD i, pos;
	PW_DBHEADER hdr;
	UINT8 uFinalKey[32];
	sha256_ctx sha32;
	USHORT usFieldType;
	DWORD dwFieldSize;
	BYTE aCompressedTime[5];

	ASSERT(pszFile != NULL);
	if(pszFile == NULL) return PWE_INVALID_PARAM;
	ASSERT(_tcslen(pszFile) != 0);
	if(_tcslen(pszFile) == 0) return PWE_INVALID_PARAM;

	_AddAllMetaStreams();

	uFileSize = sizeof(PW_DBHEADER);

	BYTE *pbt;

	// Get the size of all groups
	for(i = 0; i < m_dwNumGroups; i++)
	{
		uFileSize += 94; // 6+4+6+6+5+6+5+6+5+6+5+6+4+6+6+2+6+4 = 94

		pbt = _StringToUTF8(m_pGroups[i].pszGroupName);
		uFileSize += szlen((char *)pbt) + 1;
		SAFE_DELETE_ARRAY(pbt);
	}

	// Get the size of all entries together
	for(i = 0; i < m_dwNumEntries; i++)
	{
		ASSERT_ENTRY(&m_pEntries[i]);

		UnlockEntryPassword(&m_pEntries[i]);

		uFileSize += 134; // 6+16+6+4+6+4+6+6+6+6+6+6+5+6+5+6+5+6+5+6 = 122

		pbt = _StringToUTF8(m_pEntries[i].pszTitle);
		uFileSize += szlen((char *)pbt) + 1;
		SAFE_DELETE_ARRAY(pbt);

		pbt = _StringToUTF8(m_pEntries[i].pszUserName);
		uFileSize += szlen((char *)pbt) + 1;
		SAFE_DELETE_ARRAY(pbt);

		pbt = _StringToUTF8(m_pEntries[i].pszURL);
		uFileSize += szlen((char *)pbt) + 1;
		SAFE_DELETE_ARRAY(pbt);

		pbt = _StringToUTF8(m_pEntries[i].pszPassword);
		uFileSize += szlen((char *)pbt) + 1;
		SAFE_DELETE_ARRAY(pbt);

		pbt = _StringToUTF8(m_pEntries[i].pszAdditional);
		uFileSize += szlen((char *)pbt) + 1;
		SAFE_DELETE_ARRAY(pbt);

		pbt = _StringToUTF8(m_pEntries[i].pszBinaryDesc);
		uFileSize += szlen((char *)pbt) + 1;
		SAFE_DELETE_ARRAY(pbt);

		uFileSize += m_pEntries[i].uBinaryDataLen;

		LockEntryPassword(&m_pEntries[i]);
	}

	// Round up filesize to 16-byte boundary for Rijndael/Twofish
	uFileSize = (uFileSize + 16) - (uFileSize % 16);

	// Allocate enough memory
	uAllocated = uFileSize + 16;
	pVirtualFile = new char[uAllocated];
	ASSERT(pVirtualFile != NULL);
	if(pVirtualFile == NULL) { _LoadAndRemoveAllMetaStreams(false); return PWE_NO_MEM; }

	// Build header structure
	hdr.dwSignature1 = PWM_DBSIG_1;
	hdr.dwSignature2 = PWM_DBSIG_2;

	hdr.dwFlags = PWM_FLAG_SHA2; // The one and only hash algorithm available currently

	if(m_nAlgorithm == ALGO_AES) hdr.dwFlags |= PWM_FLAG_RIJNDAEL;
	else if(m_nAlgorithm == ALGO_TWOFISH) hdr.dwFlags |= PWM_FLAG_TWOFISH;
	else { ASSERT(FALSE); _LoadAndRemoveAllMetaStreams(false); return PWE_INVALID_PARAM; }

	hdr.dwVersion = PWM_DBVER_DW;
	hdr.dwGroups = m_dwNumGroups;
	hdr.dwEntries = m_dwNumEntries;
	hdr.dwKeyEncRounds = m_dwKeyEncRounds;

	// Make up the master key hash seed and the encryption IV
	m_random.GetRandomBuffer(hdr.aMasterSeed, 16);
	m_random.GetRandomBuffer((BYTE *)hdr.aEncryptionIV, 16);
	m_random.GetRandomBuffer(hdr.aMasterSeed2, 32);

	// Skip the header, it will be written later