Detect Driver

#pragma once

#include "IFileChecker.h"
#include "ListUtils.h"

#include <boost/bind.hpp>
#include <vector>

namespace HideAlgorithm
{

typedef NTSTATUS(*OriginalHandlerWrapperPtr)(const NtQueryDirParams& params);

struct HideParams
{
	const NtQueryDirParams& callParams;
	IFileChecker* checker;
	OriginalHandlerWrapperPtr wrapperPtr;
};

template<class InfoType>
inline bool 
CheckFileEntry( utils::NtList* list, 
			    const HideParams& params )
{
    InfoType* info = (InfoType*)list;
    return params.checker->CheckFile(info->FileName,
									 info->FileNameLength/2,
									 params.callParams);
}

inline void 
ShiftBuffer(utils::NtList* pList,utils::NtList* newBegin)
{
	size_t shiftSize = newBegin - pList;
	size_t infoSize = utils::GetListSize(pList);
	size_t newInfoSize = infoSize - shiftSize;

	std::vector<char> buffer(newInfoSize);
	memcpy(&buffer.front(),newBegin,newInfoSize);

	memcpy(pList,&buffer.front(),newInfoSize);
}

template<class InfoType>
inline NTSTATUS 
FirstEntryProcessor( const HideParams& params )
{
	utils::NtList* pList = 
		(utils::NtList*)params.callParams.FileInformation;

	utils::NtList* pCurEntry = pList;
	
	while(true)
	{
		if(!CheckFileEntry<InfoType>(pCurEntry,params))
		{
			if(pList == pCurEntry)
				break; // Nothing to hide
			
			ShiftBuffer(pList,pCurEntry);
			break; // First entry hiding complete
		}
	
		// This entry needs to be hidden
		if( utils::IsEntryLast(pCurEntry) == false )
		{
			// Move to next entry to check
			// This need to shift buffer only once
			pCurEntry = utils::GetNextEntryPointer(pCurEntry);
		}
		else
		{
			// Reached last entry 
			// This mean that all data needs to be hidden

			// Try to request more data
			NTSTATUS status = params.wrapperPtr(params.callParams);
			if( !NT_SUCCESS(status) )
				return status;
			
			// Move to begin and resume checking
			pCurEntry = pList;
		}
	}
	return STATUS_SUCCESS;
}

template<class InfoType>
inline NTSTATUS 
NextEntryProcessor( const HideParams& params )
{
	utils::NtList* pList = 
		(utils::NtList*)params.callParams.FileInformation;

	if( utils::IsEntryLast(pList) )
		return STATUS_SUCCESS;

	utils::CutFromListByFakeOffset_If(pList,    
		boost::bind(&CheckFileEntry<InfoType>,_1,params));    

	return STATUS_SUCCESS;
}

template<class InfoType>
inline NTSTATUS 
HideFileImpl( const HideParams& params )
{
	NTSTATUS status = FirstEntryProcessor<InfoType>(params);
    if( !NT_SUCCESS(status) )
		return status;

	status = NextEntryProcessor<InfoType>(params);
    if( !NT_SUCCESS(status) )
		return status;

	return STATUS_SUCCESS;
}

template<class InfoType>
inline NTSTATUS 
HideFile( const HideParams& params )
{
    try
    {
        return HideFileImpl<InfoType>(params);
    }
    catch(const std::exception& ex)
    {
#ifdef KdPrint // For use in user mode environment
		KdPrint( (__FUNCTION__" std::exception: %s\n",ex.what()) );
#endif
    }
    return STATUS_SUCCESS;
}

}