// BatchFileCmpt.cpp: implementation of the CBatchFileCmpt class.
//
//////////////////////////////////////////////////////////////////////

#include "stdafx.h"
#include "BatchFileCmpt.h"

#include "..\Common\Defines.h"
#include "resource.h"
#include <appmisc.h>

//////////////////////////////////////////////////////////////////////
// Construction/Destruction
//////////////////////////////////////////////////////////////////////

CBatchFileCmpt::CBatchFileCmpt()
{
	m_sName = _T("Finalize Shutdown");
	m_bBatchFile = false;
}

CBatchFileCmpt::~CBatchFileCmpt()
{
	if(m_bBatchFile)
		CallBatchFile();
}

//////////////////////////////////////////////////////////////////
// Calls the two helper functions below to create the batchfile
// and then run it

HRESULT CBatchFileCmpt::DoShutdown(DWORD dwMode, HWND hWndParent)
{
	// Create the batch file
	HRESULT hr = MakeBatchFile();

	if(g_site.m_log.HasErrors())
		g_site.m_log.WaitLog();

	// Calls the batch file from destructor
	m_bBatchFile = SUCCEEDED(hr) && (hr != S_FALSE) && 
				   (!g_site.m_dlgItems.IsCancelled());

	return hr;
}

//////////////////////////////////////////////////////////////////
// Creates (or modifies) the Batchfile DosMode.bat in Program Folder

HRESULT CBatchFileCmpt::MakeBatchFile()
{
	// Get Batch File Name
	string sFileName;

	sFileName.load_string(IDS_BATCHFILENAME);
	sFileName = GetProgramFolder(_Module.m_hInst) + sFileName;
	
	// Get DOS Component Strings
	string sMainBatch, sTemp;

	sTemp.load_string(IDS_STARTSECTION);
	sMainBatch += sTemp;

	HRESULT hr = S_OK;
	int nDOSComponents = 0;

	for(int nCnt = 0; nCnt < g_aComponents.size(); nCnt++)
	{
		if(g_aComponents[nCnt]->GetType() == COMPONENT_DOS)
		{
			if(g_aComponents[nCnt]->IsEnabled())
			{
				sTemp.resize(0);
				hr = g_aComponents[nCnt]->GetBatchText(sTemp);

				if(SUCCEEDED(hr)) 
				{
					sMainBatch += sTemp + "\r\n";
					nDOSComponents++;
				}

			}

		}

	}

	// If there's no DOS Components then don't do a shutdown
	if(!nDOSComponents)
		if(!_Module.m_settings.GetInt(FORCE_SHUTDOWN_KEY, false) ? true : false)
			return S_FALSE;

	sTemp.load_string(IDS_ENDSECTION);
	sMainBatch += sTemp;

	// Get any additions to the batch file;
	string sPreText;
	string sPostText;
	hr = GetCurBatchFile(sFileName, sPreText, sPostText);

	if(FAILED(hr))
		return hr;

	HANDLE fileMSModeBatch;

	// Open the File for Writing
	// This time we truncate the file
	// if it isn't already there then create it
	if((fileMSModeBatch = CreateFile(sFileName, GENERIC_WRITE,
									/*0*/FILE_SHARE_READ, NULL, OPEN_ALWAYS | TRUNCATE_EXISTING, FILE_ATTRIBUTE_NORMAL,
									NULL)) == INVALID_HANDLE_VALUE)
		return ReportLastError(_T("Couldn't write to the the batch file\n\n"));

	// Write All the differnet sections to the file
	DWORD dwWritten = 0;
	WriteFile(fileMSModeBatch, sPreText, sPreText.size(), &dwWritten, NULL);
	WriteFile(fileMSModeBatch, sMainBatch, sMainBatch.size(), &dwWritten, NULL);
	WriteFile(fileMSModeBatch, sPostText, sPostText.size(), &dwWritten, NULL);

	CloseHandle(fileMSModeBatch);

	return (nDOSComponents > 0) ? S_OK : S_FALSE;
}

HRESULT CBatchFileCmpt::GetCurBatchFile(const string& sFileName, string& sPreText, string& sPostText)
{

	// Now read in the file
	HANDLE fileMSModeBatch;

	// Open the file 
	// if it isn't already there then create it
	if((fileMSModeBatch = CreateFile(sFileName, GENERIC_READ,
									0, NULL, OPEN_ALWAYS, FILE_ATTRIBUTE_NORMAL,
									NULL)) == INVALID_HANDLE_VALUE)
		return ReportLastError(_T("Couldn't open the batch file\n\n"));

	// Load all our tags and search text
	string sStartSearchText;	sStartSearchText.load_string(IDS_STARTSEARHTEXT);
	string sEndSearchText;		sEndSearchText.load_string(IDS_ENDSEARCHTEXT);
	string sOriginal;

	string::size_type nStartPos = string::npos;
	string::size_type nEndPos = string::npos;

	// Only continue reading if there's something
	if(GetFileSize(fileMSModeBatch, NULL))
	{
		// Read Entire file in
		DWORD dwToRead = 1024;
		DWORD dwRead = 1024;

		do
		{
			dwToRead += dwToRead;
			dwRead += dwRead;

			SetFilePointer(fileMSModeBatch, 0, NULL, FILE_BEGIN);
			if(!ReadFile(fileMSModeBatch, sOriginal.get_buffer(dwToRead), dwToRead, &dwRead, NULL))
			{
				HRESULT hr = ReportLastError(_T("Couldn't read from the batch file\n\n"));
				CloseHandle(fileMSModeBatch);
				return hr;
			}

			sOriginal.release_buffer();
			sOriginal.resize(dwRead);
		}
		while(dwRead == dwToRead);


		// Get the beginning of the insert
		nStartPos = sOriginal.find(sStartSearchText);
		// Get the beginning of the line
		if(nStartPos != string::npos)
			nStartPos = sOriginal.rfind(_T("\r\n"), nStartPos);

		// We found the beginning of the line skip the \r\n
		if(nStartPos != string::npos)
			nStartPos += 2;
		// Otherwise just start from the top of the file
		else
			nStartPos = 0;

		// Get the Prefix Text
		sPreText = sOriginal.substr(0, nStartPos);



		// Get the End of the Insert
		nEndPos = sOriginal.find(sEndSearchText, nStartPos);

		// Try to find the next line
		if(nEndPos != string::npos)
			nEndPos = sOriginal.find(_T("\r\n"), nEndPos);
		// Skip over that carriage return
		if(nEndPos != string::npos)
			nEndPos += 2;

		// If we found anything then retrieve Suffix Text
		if(nEndPos != string::npos)
			sPostText = sOriginal.substr(nEndPos, string::npos);
	}

	// Now that we've read everything in close the file
	CloseHandle(fileMSModeBatch);

	return S_OK;
}

///////////////////////////////////////////////////////////////
// Runs the Batchfile Created

HRESULT CBatchFileCmpt::CallBatchFile()
{
	string sFileName;

	// Only go into DOS Mode if Windows 9x
	// The different PIFs take care of that
	if(::GetPlatform() == VER_PLATFORM_WIN32_NT)
		sFileName.load_string(IDS_BATCHFILEPIF_NT);
	else
		sFileName.load_string(IDS_BATCHFILEPIF);


	// Make Batchfile Name
	string sTemp;
	string sFileDir = GetProgramFolder(_Module.m_hInst);

	sFileName = sFileDir + sFileName;

#ifndef _DEBUG

	// Run the batchfile (PIF)
	if(ShellExecute(NULL, _T("open"), sFileName, NULL, sFileDir, SW_SHOWNORMAL) <= (HINSTANCE)32)
		::MessageBox(NULL, _T("Couldn't start MS-DOS mode batch file."), _T("Secure Shutdown"), MB_ICONSTOP);

#endif // _DEBUG

	return S_OK;
}



HRESULT CBatchFileCmpt::ReportLastError(string sMessage)
{
	HRESULT hr = HRESULT_FROM_WIN32(::GetLastError());
	sMessage += FormatHR(hr);
	MessageBox(NULL, sMessage, _T("Secure Shutdown"), MB_ICONSTOP | MB_OK);
	return hr;
}