//
// $Id: FileFinder.cpp,v 1.8 2006/08/18 14:01:45 bakerj Exp $
//
//****************************************************************************************//
// Copyright (c) 2006, 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  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
void FileFinder::FindPaths(string regex, StringVector* paths, bool isRegex) {
	// -----------------------------------------------------------------------
	//
	//  ABSTRACT
	//
	//	Search the file system for paths that match the provided regex.
	//	Attempt to break off a constant portion of the path to spead up searching
	// -----------------------------------------------------------------------

	StringVector* drives;
	string fileName = "";

	string patternOut= "";
	string constPortion= "";
	string fileSeperatorStr = "";
	fileSeperatorStr+=Common::fileSeperator;
	this->fileMatcher->GetConstantPortion(regex, fileSeperatorStr, &patternOut, &constPortion);
	// Remove extra slashes
	constPortion = this->fileMatcher->RemoveExtraSlashes(constPortion);

	// Found a constant portion
	if(constPortion.compare("") != 0 && patternOut.compare("") != 0) {

		//	Call search function
		this->GetPathsForPattern(constPortion, regex, paths);

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

		StringVector::iterator drive;
		for (drive=drives->begin(); drive!=drives->end(); drive++) {
			//	Call search function
			try  {

				this->GetPathsForPattern((*drive), regex, paths, isRegex);

			} catch(REGEXException ex) {
				if(ex.GetSeverity() == ERROR_WARN) {
					string pcreMsg = "";
					pcreMsg.append("Filefinder Warning - while searching for matching files:\n");
					pcreMsg.append("-----------------------------------------------------------------------\n");
					pcreMsg.append(ex.GetErrorMessage());
					Log::Debug(pcreMsg);
				} else {
					throw;
				}		
			}
		}
		delete drives;

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

		//	There are no pattern matching chars treat this as a normal path 
		if(this->PathExists(constPortion)) {
			paths->push_back(constPortion);
		}
	}

}

StringVector* FileFinder::GetDrives() {
	// -----------------------------------------------------------------------
	//
	//  ABSTRACT
	//
	//	Get all fixed drives on the system. Return them in a string vector
	// -----------------------------------------------------------------------
	StringVector* drives = new StringVector();	
	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");
		throw FileFinderException(errMsg);

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

	
	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");
		throw FileFinderException(errMsg);

	//	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::GetPathsForPattern(string dirIn, string pattern, StringVector *pathVector, bool isRegex) {
	// -----------------------------------------------------------------------
	//
	//  ABSTRACT
	//
	//  This function gets all paths that match a given pattern.
	//	It will not attempt to match any 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] != Common::fileSeperator)
			dirIn.append(1, Common::fileSeperator);
		
		// 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 = INVALID_HANDLE_VALUE;

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

			string errorMessage = "";
			errorMessage.append("Error: Unable to get a valid handle in GetPathsForPattern().\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);

				if(this->IsMatch(pattern, dirToSearch, isRegex))
					pathVector->push_back(dirToSearch);

				GetPathsForPattern(dirToSearch, pattern, pathVector);
			
			}
		} 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 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);
	}
}

void FileFinder::GetFilesForPattern(string path, string pattern, BehaviorVector* behaviors, StringVector* fileNames, bool isRegex) {
	// -----------------------------------------------------------------------
	//
	//  ABSTRACT
	//
	//  This function gets all file names that match a given pattern on the spaecified path.
	//	If a match is found the filename is pushed on to a vector of strings.
	//
	// -----------------------------------------------------------------------

	try {

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

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

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

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

		//	Loop through each file in the directory.  
		//	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) {

				// skip
			
			} else {
			
				string fileName = FindFileData.cFileName;
	
				//	Check pattern
				if(this->IsMatch(pattern, fileName, isRegex))
					fileNames->push_back(fileName);
			}
		} 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(path);
			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(path);
		errorMessage.append("\n\tPattern: ");
		errorMessage.append(pattern);
		throw FileFinderException(errorMessage);
	}
}

bool FileFinder::PathExists(string path) {
	// -----------------------------------------------------------------------
	//
	//  ABSTRACT
	//
	//  Return true if the specified path exists.
	//
	// -----------------------------------------------------------------------

	bool exists = false;

	if (path[path.length()-1] != Common::fileSeperator)
		path.append(1, Common::fileSeperator);


	HANDLE hFile = INVALID_HANDLE_VALUE;

	try {
		hFile = CreateFile(path.c_str(),					// DirName
								GENERIC_READ,				// access mode
								FILE_SHARE_READ,			// share mode
								NULL,						// SD
								OPEN_EXISTING,				// how to create
								FILE_FLAG_BACKUP_SEMANTICS,	// file attributes
								NULL);						// handle to template file


		if (hFile == INVALID_HANDLE_VALUE) {

			DWORD errorNum = GetLastError();

			if(errorNum == ERROR_FILE_NOT_FOUND) {
				// skip this not looking for a file
				
			} else if(errorNum == ERROR_PATH_NOT_FOUND) {

				// skip this the path does not exist

			} else {
				char errorCodeBuffer[33];
				_ltoa(errorNum, errorCodeBuffer, 10);

				string errorMessage = "";
				errorMessage.append("(FileProbe) Unable to open a handle to the file '");
				errorMessage.append(path);
				errorMessage.append("'.  Error Code - ");
				errorMessage.append(errorCodeBuffer);
				throw FileFinderException(errorMessage);
			}
		} else {
			exists = true;
		}
		CloseHandle(hFile); 

	} catch(Exception ex) {
		CloseHandle(hFile); 
		throw;

	} catch(...) {
		CloseHandle(hFile); 
		string errorMessage = "";
		errorMessage.append("Error: ");
		errorMessage.append("An unspecified error was encountered while trying to search for matching paths. \n\tDirectory: ");
		errorMessage.append(path);
		throw FileFinderException(errorMessage);
	}

	return exists;
}

bool FileFinder::FileNameExists(string path, string fileName) {
	// -----------------------------------------------------------------------
	//
	//  ABSTRACT
	//
	//  Return true if the specified file exists on the specifeid path.
	//
	// -----------------------------------------------------------------------

	bool exists = false;

	if (path[path.length()-1] != Common::fileSeperator)
		path.append(1, Common::fileSeperator);

	string filePath = path + fileName;

	HANDLE hFile = INVALID_HANDLE_VALUE;

	try {
		hFile = CreateFile(filePath.c_str(),				// file name
								GENERIC_READ,				// access mode
								FILE_SHARE_READ,			// share mode
								NULL,						// SD
								OPEN_EXISTING,				// how to create
								FILE_ATTRIBUTE_NORMAL,		// file attributes
								NULL);						// handle to template file

		if (hFile == INVALID_HANDLE_VALUE) {

			DWORD errorNum = GetLastError();

			if(errorNum == ERROR_FILE_NOT_FOUND) {
				// skip this the file does not exist
				
			} else if(errorNum == ERROR_PATH_NOT_FOUND) {

				// skip this since if the path is not found the file can't exist

			} else {
				char errorCodeBuffer[33];
				_ltoa(errorNum, errorCodeBuffer, 10);

				string errorMessage = "";
				errorMessage.append("(FileProbe) Unable to open a handle to the file '");
				errorMessage.append(filePath);
				errorMessage.append("'.  Error Code - ");
				errorMessage.append(errorCodeBuffer);
				throw FileFinderException(errorMessage);
			}
		} else {
			exists = true;
		}
		CloseHandle(hFile); 

	} catch(Exception ex) {
		CloseHandle(hFile); 
		throw;

	} catch(...) {
		CloseHandle(hFile); 
		string errorMessage = "";
		errorMessage.append("Error: ");
		errorMessage.append("An unspecified error was encountered while trying to search for matching paths. \n\tDirectory: ");
		errorMessage.append(path);
		throw FileFinderException(errorMessage);
	}

	return exists;
}