archiverecovery.cpp

非常出名开源客户端下载的程序emule

#include "stdafx.h"
#include "ArchiveRecovery.h"

#ifdef _DEBUG
#undef THIS_FILE
static char THIS_FILE[]=__FILE__;
#define new DEBUG_NEW
#endif


// At some point it may be worth displaying messages to alert the user if there were errors, or where to find the file.

void CArchiveRecovery::recover(CPartFile *partFile, bool preview)
{
	if (partFile->m_bPreviewing || partFile->m_bRecoveringArchive)
		return;
	partFile->m_bRecoveringArchive = true;

	theApp.emuledlg->AddLogLine(true, GetResString(IDS_ATTEMPTING_RECOVERY) );

	// Get the current filled list for this file
	CTypedPtrList<CPtrList, Gap_Struct*> *filled = new CTypedPtrList<CPtrList, Gap_Struct*>;
	partFile->GetFilledList(filled);

	// The rest of the work can be safely done in a new thread
	ThreadParam *tp = new ThreadParam;
	tp->partFile = partFile;
	tp->filled = filled;
	tp->preview = preview;
	DWORD dwThreadId;
	HANDLE hThread = CreateThread(	NULL,			// no security attributes 
									0,				// use default stack size  
									run,			// thread function 
									tp,				// argument to thread function 
									0,				// use default creation flags 
									&dwThreadId);	// returns the thread identifier 

	// Check the return value for success. 
	if (hThread == NULL) 
	{
		partFile->m_bRecoveringArchive = false;
		theApp.emuledlg->AddLogLine(true, GetResString(IDS_RECOVERY_FAILED));
		// Need to delete the memory here as won't be done in thread
		DeleteMemory(tp);
	}
}

DWORD WINAPI CArchiveRecovery::run(LPVOID lpParam)
{
	ThreadParam *tp = (ThreadParam *)lpParam;

	if (!performRecovery(tp->partFile, tp->filled, tp->preview))
		theApp.emuledlg->AddLogLine(true, GetResString(IDS_RECOVERY_FAILED));

	tp->partFile->m_bRecoveringArchive = false;

	// Delete memory used by copied gap list
	DeleteMemory(tp);

	return 0;
}

bool CArchiveRecovery::performRecovery(CPartFile *partFile, CTypedPtrList<CPtrList, Gap_Struct*> *filled, bool preview)
{
	bool success = false;
	try
	{
		// Copy the file
		CString tempFileName = CString(theApp.glob_prefs->GetTempDir()) + CString("\\") + CString(partFile->GetFileName()).Mid(0,5) + CString("-rec.tmp");
		if (!CopyFile(partFile, filled, tempFileName))
			return false;

		// Open temp file for reading
		CFile temp;
		if (!temp.Open(tempFileName, CFile::modeRead))
			return false;

		// Open the output file
		CString ext = CString(partFile->GetFileName()).Right(4);
		CString outputFileName = CString(theApp.glob_prefs->GetTempDir()) + CString("\\") + CString(partFile->GetFileName()).Mid(0,5) + CString("-rec") + ext;
		CFile output;
		if (output.Open(outputFileName, CFile::modeWrite | CFile::shareExclusive | CFile::modeCreate))
		{
			// Process the output file
			if (ext.CompareNoCase(".zip") == 0)
				success = recoverZip(&temp, &output, filled, (temp.GetLength() == partFile->GetFileSize()));
			else if (ext.CompareNoCase(".rar") == 0)
				success = recoverRar(&temp, &output, filled);

			// Close output
			output.Close();
		}
		// Close temp file
		temp.Close();

		// Report success
		if (success)
		{
			theApp.emuledlg->AddLogLine(true, GetResString(IDS_RECOVERY_SUCCESSFUL));
			// Preview file if required
			if (preview)
			{
				SHELLEXECUTEINFO SE;
				memset(&SE,0,sizeof(SE));
				SE.fMask = SEE_MASK_NOCLOSEPROCESS ;
				SE.lpVerb = "open";
				SE.lpFile = outputFileName.GetBuffer();
				SE.nShow = SW_SHOW;
				SE.cbSize = sizeof(SE);
				ShellExecuteEx(&SE);
				if (SE.hProcess)
				{
					WaitForSingleObject(SE.hProcess, INFINITE);
					CloseHandle(SE.hProcess);
				}
				CFile::Remove(outputFileName);
			}
		}

		// Remove temp file
		CFile::Remove(tempFileName);
	} catch (...) {}
	return success;
}

bool CArchiveRecovery::recoverZip(CFile *zipInput, CFile *zipOutput, CTypedPtrList<CPtrList, Gap_Struct*> *filled, bool fullSize)
{
	bool retVal = false;
	long fileCount = 0;
	try
	{
		CTypedPtrList<CPtrList, ZIP_CentralDirectory*> centralDirectoryEntries;
		Gap_Struct *fill;

		// If the central directory is intact this is simple
		if (fullSize && readZipCentralDirectory(zipInput, &centralDirectoryEntries, filled))
		{
			if (centralDirectoryEntries.GetCount() == 0)
				return false;
			ZIP_CentralDirectory *cdEntry;
			POSITION pos = centralDirectoryEntries.GetHeadPosition();
			bool deleteCD;
			for (int i=centralDirectoryEntries.GetCount(); i>0; i--)
			{
				deleteCD = false;
				cdEntry = centralDirectoryEntries.GetAt(pos);
				uint32 lenEntry = sizeof(ZIP_Entry) + cdEntry->lenFilename + cdEntry->lenExtraField + cdEntry->lenCompressed;
				if (IsFilled(cdEntry->relativeOffsetOfLocalHeader, cdEntry->relativeOffsetOfLocalHeader + lenEntry, filled))
				{
					zipInput->Seek(cdEntry->relativeOffsetOfLocalHeader, CFile::begin);
					// Update offset
					cdEntry->relativeOffsetOfLocalHeader = zipOutput->GetPosition();
					if (!processZipEntry(zipInput, zipOutput, lenEntry, NULL))
						deleteCD = true;
				}
				else
					deleteCD = true;

				if (deleteCD)
				{
					delete [] cdEntry->filename;
					if (cdEntry->lenExtraField > 0)
						delete [] cdEntry->extraField;
					if (cdEntry->lenComment > 0)
						delete [] cdEntry->comment;
					delete cdEntry;
					POSITION del = pos;
					centralDirectoryEntries.GetNext(pos);
					centralDirectoryEntries.RemoveAt(del);
				}
				else
					centralDirectoryEntries.GetNext(pos);
			}
		}
		else // Have to scan the file the hard way
		{		
			// Loop through filled areas of the file looking for entries
			POSITION pos = filled->GetHeadPosition();
			while (pos != NULL)
			{
				fill = filled->GetNext(pos);
				uint32 filePos = zipInput->GetPosition();
				// The file may have been positioned to the next entry in ScanForMarker() or processZipEntry()
				if (filePos > fill->end)
					continue;
				if (filePos < fill->start)
					zipInput->Seek(fill->start, CFile::begin);

				// If there is any problem, then don't bother checking the rest of this part
				while (true)
				{
					// Scan for entry marker within this filled area
					if (!scanForZipMarker(zipInput, (uint32)ZIP_LOCAL_HEADER_MAGIC, (fill->end - zipInput->GetPosition() + 1)))
						break;
					if (zipInput->GetPosition() > fill->end)
						break;
					if (!processZipEntry(zipInput, zipOutput, (fill->end - zipInput->GetPosition() + 1), &centralDirectoryEntries))
						break;
				}
			}
		}

		// Remember offset before CD entries
		uint32 startOffset = zipOutput->GetPosition();

		// Write all central directory entries
		fileCount = centralDirectoryEntries.GetCount();
		if (fileCount > 0)
		{
			ZIP_CentralDirectory *cdEntry;
			POSITION pos = centralDirectoryEntries.GetHeadPosition();
			while (pos != NULL)
			{
				cdEntry = centralDirectoryEntries.GetNext(pos);

				writeUInt32(zipOutput, ZIP_CD_MAGIC);
				writeUInt16(zipOutput, cdEntry->versionMadeBy);
				writeUInt16(zipOutput, cdEntry->versionToExtract);
				writeUInt16(zipOutput, cdEntry->generalPurposeFlag);
				writeUInt16(zipOutput, cdEntry->compressionMethod);
				writeUInt16(zipOutput, cdEntry->lastModFileTime);
				writeUInt16(zipOutput, cdEntry->lastModFileDate);
				writeUInt32(zipOutput, cdEntry->crc32);
				writeUInt32(zipOutput, cdEntry->lenCompressed);
				writeUInt32(zipOutput, cdEntry->lenUnompressed);
				writeUInt16(zipOutput, cdEntry->lenFilename);
				writeUInt16(zipOutput, cdEntry->lenExtraField);
				writeUInt16(zipOutput, cdEntry->lenComment);
				writeUInt16(zipOutput, 0); // Disk number start
				writeUInt16(zipOutput, cdEntry->internalFileAttributes);
				writeUInt32(zipOutput, cdEntry->externalFileAttributes);
				writeUInt32(zipOutput, cdEntry->relativeOffsetOfLocalHeader);
				zipOutput->Write(cdEntry->filename, cdEntry->lenFilename);
				if (cdEntry->lenExtraField > 0)
					zipOutput->Write(cdEntry->extraField, cdEntry->lenExtraField);
				if (cdEntry->lenComment > 0)
					zipOutput->Write(cdEntry->comment, cdEntry->lenComment);

				delete [] cdEntry->filename;
				if (cdEntry->lenExtraField > 0)
					delete [] cdEntry->extraField;
				if (cdEntry->lenComment > 0)
					delete [] cdEntry->comment;
				delete cdEntry;
			}

			// Remember offset before CD entries
			uint32 endOffset = zipOutput->GetPosition();

			// Write end of central directory
			writeUInt32(zipOutput, ZIP_END_CD_MAGIC);
			writeUInt16(zipOutput, 0); // Number of this disk
			writeUInt16(zipOutput, 0); // Number of the disk with the start of the central directory
			writeUInt16(zipOutput, fileCount);
			writeUInt16(zipOutput, fileCount);
			writeUInt32(zipOutput, endOffset - startOffset);
			writeUInt32(zipOutput, startOffset);
			writeUInt16(zipOutput, strlen(ZIP_COMMENT));
			zipOutput->Write(ZIP_COMMENT, strlen(ZIP_COMMENT));

			centralDirectoryEntries.RemoveAll();
		}
		retVal = true;
	} catch (...) {}

	// Tell the user how many files were recovered
	CString msg;	
	if (fileCount == 1)
		msg = GetResString(IDS_RECOVER_SINGLE);
	else
		msg.Format(GetResString(IDS_RECOVER_MULTIPLE), fileCount);
	//AfxMessageBox(msg, MB_OK | MB_ICONINFORMATION);
	theApp.emuledlg->AddLogLine(true,msg);

	return retVal;
}

bool CArchiveRecovery::readZipCentralDirectory(CFile *zipInput, CTypedPtrList<CPtrList, ZIP_CentralDirectory*> *centralDirectoryEntries, CTypedPtrList<CPtrList, Gap_Struct*> *filled)
{
	bool retVal = false;
	try
	{
		// Ideally this zip file will not have a comment and the End-CD will be easy to find
		zipInput->Seek(-22, CFile::end);
		if (!(readUInt32(zipInput) == ZIP_END_CD_MAGIC))
		{
			// Have to look for it, comment could be up to 65535 chars but only try with less than 1k
			zipInput->Seek(-1046, CFile::end);
			if (!scanForZipMarker(zipInput, (uint32)ZIP_END_CD_MAGIC, 1046))
				return false;
			// Skip it again
			readUInt32(zipInput);
		}

		// Found End-CD
		// Only interested in offset of first CD
		zipInput->Seek(12, CFile::current);
		uint32 startOffset = readUInt32(zipInput);
		if (!IsFilled(startOffset, zipInput->GetLength(), filled))
			return false;

		// Goto first CD and start reading
		zipInput->Seek(startOffset, CFile::begin);
		ZIP_CentralDirectory *cdEntry;
		while (readUInt32(zipInput) == ZIP_CD_MAGIC)
		{
			cdEntry = new ZIP_CentralDirectory;
			cdEntry->versionMadeBy				= readUInt16(zipInput);
			cdEntry->versionToExtract			= readUInt16(zipInput);
			cdEntry->generalPurposeFlag			= readUInt16(zipInput);
			cdEntry->compressionMethod			= readUInt16(zipInput);
			cdEntry->lastModFileTime			= readUInt16(zipInput);
			cdEntry->lastModFileDate			= readUInt16(zipInput);
			cdEntry->crc32						= readUInt32(zipInput);
			cdEntry->lenCompressed				= readUInt32(zipInput);
			cdEntry->lenUnompressed				= readUInt32(zipInput);
			cdEntry->lenFilename				= readUInt16(zipInput);
			cdEntry->lenExtraField				= readUInt16(zipInput);
			cdEntry->lenComment					= readUInt16(zipInput);
			cdEntry->diskNumberStart			= readUInt16(zipInput);
			cdEntry->internalFileAttributes		= readUInt16(zipInput);
			cdEntry->externalFileAttributes		= readUInt32(zipInput);
			cdEntry->relativeOffsetOfLocalHeader= readUInt32(zipInput);

			if (cdEntry->lenFilename > 0)
			{
				cdEntry->filename					= new BYTE[cdEntry->lenFilename];
				zipInput->Read(cdEntry->filename, cdEntry->lenFilename);
			}
			if (cdEntry->lenExtraField > 0)
			{
				cdEntry->extraField					= new BYTE[cdEntry->lenExtraField];
				zipInput->Read(cdEntry->extraField, cdEntry->lenExtraField);
			}
			if (cdEntry->lenComment > 0)
			{
				cdEntry->comment					= new BYTE[cdEntry->lenComment];
				zipInput->Read(cdEntry->comment, cdEntry->lenComment);
			}

			centralDirectoryEntries->AddTail(cdEntry);
		}

		retVal = true;
	} catch (...) {}
	return retVal;
}

bool CArchiveRecovery::processZipEntry(CFile *zipInput, CFile *zipOutput, uint32 available, CTypedPtrList<CPtrList, ZIP_CentralDirectory*> *centralDirectoryEntries)
{
	if (available < 26)
		return false;

	bool retVal = false;
	try
	{
		// Need to know where it started
		long startOffset = zipOutput->GetPosition();

		// Entry format :
		//  4      2 bytes  Version needed to extract
		//  6      2 bytes  General purpose bit flag
		//  8      2 bytes  Compression method
		// 10      2 bytes  Last mod file time
		// 12      2 bytes  Last mod file date
		// 14      4 bytes  CRC-32
		// 18      4 bytes  Compressed size (n)
		// 22      4 bytes  Uncompressed size
		// 26      2 bytes  Filename length (f)
		// 28      2 bytes  Extra field length (e)
		//        (f)bytes  Filename
		//        (e)bytes  Extra field
		//        (n)bytes  Compressed data

		// Read header
		if (readUInt32(zipInput) != ZIP_LOCAL_HEADER_MAGIC)
			return false;

		ZIP_Entry entry;
		entry.versionToExtract		= readUInt16(zipInput);
		entry.generalPurposeFlag	= readUInt16(zipInput);
		entry.compressionMethod		= readUInt16(zipInput);
		entry.lastModFileTime		= readUInt16(zipInput);
		entry.lastModFileDate		= readUInt16(zipInput);
		entry.crc32					= readUInt32(zipInput);
		entry.lenCompressed			= readUInt32(zipInput);
		entry.lenUncompressed		= readUInt32(zipInput);
		entry.lenFilename			= readUInt16(zipInput);
		entry.lenExtraField			= readUInt16(zipInput);
		
		// Do some quick checks at this stage that data is looking ok
		if ((entry.crc32 == 0) || (entry.lenCompressed == 0) || (entry.lenUncompressed == 0) || (entry.lenFilename == 0))
			return false;

		// Is this entry complete
		if ((entry.lenFilename + entry.lenExtraField + entry.lenCompressed) > (available - 26))
		{
			// Move the file pointer to the start of the next entry
			zipInput->Seek((entry.lenFilename + entry.lenExtraField + entry.lenCompressed), CFile::current);
			return false;
		}

		// Filename
		if (entry.lenFilename > MAX_PATH)
			return false; // Possibly corrupt, don't allocate lots of memory
		entry.filename = new BYTE[entry.lenFilename];
		if (zipInput->Read(entry.filename, entry.lenFilename) != entry.lenFilename)
		{
			delete [] entry.filename;
			return false;
		}

		// Extra data
		if (entry.lenExtraField > 0)
		{
			entry.extraField = new BYTE[entry.lenExtraField];
			zipInput->Read(entry.extraField, entry.lenExtraField);
		}

		// Output
		writeUInt32(zipOutput, ZIP_LOCAL_HEADER_MAGIC);
		writeUInt16(zipOutput, entry.versionToExtract);
		writeUInt16(zipOutput, entry.generalPurposeFlag);
		writeUInt16(zipOutput, entry.compressionMethod);
		writeUInt16(zipOutput, entry.lastModFileTime);
		writeUInt16(zipOutput, entry.lastModFileDate);
		writeUInt32(zipOutput, entry.crc32);
		writeUInt32(zipOutput, entry.lenCompressed);
		writeUInt32(zipOutput, entry.lenUncompressed);
		writeUInt16(zipOutput, entry.lenFilename);
		writeUInt16(zipOutput, entry.lenExtraField);
		if (entry.lenFilename > 0)
			zipOutput->Write(entry.filename, entry.lenFilename);
		if (entry.lenExtraField > 0)
			zipOutput->Write(entry.extraField, entry.lenExtraField);

		// Read and write compressed data to avoid reading all into memory
		uint32 written = 0;
		BYTE buf[4096];
		uint32 lenChunk = 4096;
		while (written < entry.lenCompressed)
		{
			lenChunk  = (entry.lenCompressed - written);
			if (lenChunk > 4096)
				lenChunk = 4096;
			lenChunk = zipInput->Read(buf, lenChunk);
			if (lenChunk == 0)
				break;
			written += lenChunk;
			zipOutput->Write(buf, lenChunk);
		}