//
// $Id: FileFinder.cpp,v 1.3 2005/04/11 21:02:28 bakerj Exp $
//
//****************************************************************************************//
// Copyright (c) 2005, The MITRE Corporation
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without modification, are
// permitted provided that the following conditions are met:
//
//     * Redistributions of source code must retain the above copyright notice, this list
//       of conditions and the following disclaimer.
//     * Redistributions in binary form must reproduce the above copyright notice, this 
//       list of conditions and the following disclaimer in the documentation and/or other
//       materials provided with the distribution.
//     * Neither the name of The MITRE Corporation nor the names of its contributors may be
//       used to endorse or promote products derived from this software without specific 
//       prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY 
// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
// OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT 
// SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT
// OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
// HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR
// TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE,
// EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
//
//****************************************************************************************//

#include "FileFinder.h"


FileFinder::FileFinder()
{
	// -----------------------------------------------------------------------
	//
	//  ABSTRACT
	//
	//	Note that this is a protected constructor
	// -----------------------------------------------------------------------

}

FileFinder::~FileFinder()
{
	// -----------------------------------------------------------------------
	//
	//  ABSTRACT
	//
	//	Do nothing for now
	// -----------------------------------------------------------------------

}

//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~  Private Members  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
sVector FileFinder::FindFiles(string regex)
{
	// -----------------------------------------------------------------------
	//
	//  ABSTRACT
	//
	//	Search the file system for files whose path matches the provided regex.
	//  Return a vector of matching file paths.  
	//	Attempt to break off a constant portion of the path to spead up searching
	// -----------------------------------------------------------------------

	sVector paths;
	sVector drives;
	string fileName = "";

	string patternOut= "";
	string constPortion= "";
	this->fileMatcher->GetConstantPortion(regex, "\\", &patternOut, &constPortion);
	// Remove extra slashes
	constPortion = this->fileMatcher->RemoveExtraSlashes(constPortion);

	// Found a constant porion
	if(patternOut.compare("") != 0) {

		//	Call search function
		GetFilePathsForPattern(constPortion, regex, &paths);

	//	No constant portion.
	} else if(constPortion.compare("") == 0) { 
		
		//	Must start searching for matches on each drive.
		drives = GetDrives();

		sVector::iterator drive;
		for (drive=drives.begin(); drive!=drives.end(); drive++) {
			//	Call search function
			try  {

				GetFilePathsForPattern((*drive), regex, &paths);

			} catch(REGEXException ex) {
				if(ex.GetSeverity() == ERROR_WARN) {
					string pcreMsg = "";
					pcreMsg.append("File Finder Warning:\n");
					pcreMsg.append("------------------------------\n");
					pcreMsg.append(ex.GetErrorMessage());
					if(Log::verboseMode) {
						cout << pcreMsg << "\n" << endl;
						Log::WriteLog(pcreMsg + "\n\n");
					}
				}
				break;			
			}
		}

	} else if(patternOut.compare("") == 0) {

		//	There are no pattern matching chars treat this as a normal path 
		paths.push_back(constPortion);
	}

	return paths;
}

sVector FileFinder::GetDrives()
{
	// -----------------------------------------------------------------------
	//
	//  ABSTRACT
	//
	//	Get all fixed drives on the system. Return them in a string vector
	// -----------------------------------------------------------------------
	sVector drives;	
	unsigned int index	= 0;
	string tmp			= "";
	string drive		= "";
	string errMsg		= "";
	DWORD nBufferLength = 0;
	DWORD dwResult		= 0;
	LPTSTR lpBuffer		= new char[0];

	//	Get the required buffer size
	dwResult = GetLogicalDriveStrings(	nBufferLength,	// size of buffer
										lpBuffer);		// drive strings buffer

	if(dwResult > nBufferLength) {
		//	Call the function again with the correct buffer size
		delete [] (lpBuffer); 
		lpBuffer = new char[dwResult]; 
		nBufferLength = dwResult;
		dwResult = GetLogicalDriveStrings(	nBufferLength,	// size of buffer
											lpBuffer);		// drive strings buffer
	
	} else if(dwResult == 0) {
		//	Error check GetLastError 
		char strErrorCode[33];
		_itoa(GetLastError(), strErrorCode, 10);
		errMsg.append("Error: Unable to enumerate the drives on the system. Error code: ");
		errMsg.append(strErrorCode);
		errMsg.append("\n");

	} else {
		//	Unknown Error
		errMsg.append("Error: Unable to enumerate the drives on the system. (Unknown error)\n");
	}

	
	if(dwResult == 0) {

		//	Error check GetLastError 
		char strErrorCode[33];
		_itoa(GetLastError(), strErrorCode, 10);
		errMsg.append("Error: Unable to enumerate the drives on the system. Error code: ");
		errMsg.append(strErrorCode);
		errMsg.append("\n");

	//	Process the list of drives
	} else {
		while(index < dwResult) {

			tmp = lpBuffer[index];
			index += 4;
			drive.append(tmp);
			drive.append(":\\");
			
			//	Only fixed drives
			if(GetDriveType(drive.c_str()) == DRIVE_FIXED)
				drives.push_back(drive);

			drive = "";			
		}	
	}

	return drives;
}

void FileFinder::GetFilePathsForPattern(string dirIn, string pattern, sVector *pathVector)
{
	// -----------------------------------------------------------------------
	//
	//  ABSTRACT
	//
	//  This function gets all file paths that match a given pattern.
	//	It will not attempt to match any file or directory that starts with a period or
	//	is named "System Volume Information". This is to avoid doing anything to the "." 
	//	and ".." files or accessing restricted direcoties.  
	//	This does call itself recursively as it must search all sub directories of dirIn.
	//	If a match is found the path is pushed on to a vector of strings.
	//
	// -----------------------------------------------------------------------

	try {

		//	Stop is a Null Dir
		if ((dirIn.empty() == true) || (dirIn == ""))
			return;

		// Verify that the path that was passed into this function ends with a slash.  If
		// it doesn't, then add one.
		if (dirIn[dirIn.length()-1] != '\\')
			dirIn.append(1, '\\');
		
		// Append a '*' to the end of the path to signify that we want to find all files
		// in given directory.
		string findDir;
		findDir = dirIn + "*";

		// Find the first file in the directory.  If this fails, then there is no reason
		// to continue.
		WIN32_FIND_DATA FindFileData;
		HANDLE hFind;

		hFind = FindFirstFile(findDir.c_str(), &FindFileData);
		if (hFind == INVALID_HANDLE_VALUE) {

			string errorMessage = "";
			errorMessage.append("Error: Unable to get a valid handle in GetFilePathsForPattern().\n\tDirectory: ");
			errorMessage.append(dirIn);
			errorMessage.append("\n\tPattern: ");
			errorMessage.append(pattern);
			throw FileFinderException(errorMessage);
		}

		//	Loop through each file in the directory.  
		//	If a sub-directory is found, make a recursive call to GetFilePathsForPattern to search its contents.
		//	If a file is found get the file path and check it against the pattern

		do {

			// Skip ., .., and System Volume 
			if ((strncmp(FindFileData.cFileName, ".", 1) == 0) ||
				(strncmp(FindFileData.cFileName, "..", 2) == 0) ||
				(strncmp(FindFileData.cFileName, "System Volume Information", 25) == 0))
			{
			
			// Found a dir
			} else if (FindFileData.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) {

				string dirToSearch = dirIn;					
				dirToSearch.append(FindFileData.cFileName);

				GetFilePathsForPattern(dirToSearch, pattern, pathVector);
			
			} else {

				string filePath = dirIn;					
				filePath.append(FindFileData.cFileName);
	
				//	Check pattern
				if(this->fileMatcher->IsMatch(pattern.c_str(), filePath.c_str()))
					pathVector->push_back(filePath);
			}
		} while (FindNextFile(hFind, &FindFileData));

		//	Close the handle to the file search object.
		if(!FindClose(hFind)) {
			string errorMessage = "";
			errorMessage.append("Error: ");
			errorMessage.append("Unable to close search handle while trying to search for matching file paths. \n\tDirectory: ");
			errorMessage.append(dirIn);
			errorMessage.append("\n\tPattern: ");
			errorMessage.append(pattern);

			throw FileFinderException(errorMessage);	
		}

	//	Just need to ensure that all exceptions have a nice message. 
	//	So rethrow the exceptions I created catch the others and format them.
	} catch(Exception ex) {

		throw;

	} catch(...) {

		string errorMessage = "";
		errorMessage.append("Error: ");
		errorMessage.append("An unspecified error was encountered while trying to search for matching paths. \n\tDirectory: ");
		errorMessage.append(dirIn);
		errorMessage.append("\n\tPattern: ");
		errorMessage.append(pattern);
		throw FileFinderException(errorMessage);
	}
}