#include "stdafx.h"
#include "MasterImgGen.h"
#include "BootParamDef.h"
#include <io.h>

#define SECTOR_SIZE			512
static const unsigned char MAGIC[] = { 'S', 'T', 'B', 'P' };

#define NEW_ARRAY( p, type, size ) \
{                                  \
	(p) = new type [ size ];       \
	if ( NULL == (p) )             \
		__leave;                   \
}

#define DELETE_ARRAY( p ) \
	if ( NULL != (p) ) {  \
		delete[] (p);     \
		(p) = NULL;       \
	} 

#define CLOSE_FILE( f )  \
	if ( NULL != (f) ) { \
		::fclose( f );   \
		(f) = NULL;      \
	}



void *CMasterImgGen::MakeMasterImage(DWORD *pImageSize,
									 int NumOfImages,
									 IMAGE_PARAM * lpImageParam,
									 int PageType)
{
	ASSERT( NULL != pImageSize );
	ASSERT( NumOfImages > 0 );
	ASSERT( NULL != lpImageParam );
	ASSERT((SMALL_PAGE == PageType) || (LARGE_PAGE == PageType));

	void * pImage = NULL;
	FILE * pMaster = NULL;
	m_SctsPerBlock = (SMALL_PAGE == PageType) ? 32 : 256;

	TCHAR szPath[MAX_PATH] = {0};

	/* @hongliang.xin 2009-9-2 modify the temp file location */
	LPTSTR lpszPath = _tgetenv(_T("tmp"));
	_stprintf(szPath,_T("%s\\tmp%d"),lpszPath,GetTickCount());
	
	__try {

		/* @hongliang.xin 2009-9-2 modify the temp file location */
		pMaster = ::_tfopen(szPath,_T("w+b"));

		//noticed by hongliang.xin 2009-9-2
		//pMaster = ::tmpfile();

		if ( NULL == pMaster )
			__leave;

		if ( !CalcImageOffset( NumOfImages, lpImageParam ) )
			__leave;

		if (!AddBootParam( pMaster, NumOfImages, lpImageParam ) )
			__leave;

		for ( int i=0; i<NumOfImages; ++i )
			if ( !MakeImage( pMaster, lpImageParam + i ) )
				__leave;

		VERIFY( 0 == ::fflush( pMaster ) );
		*pImageSize = ::_filelength( _fileno( pMaster ) );
		pImage = new unsigned char[*pImageSize];
		if (NULL == pImage)
			__leave;

		VERIFY(0 == ::fseek(pMaster, 0, SEEK_SET));
		VERIFY(*pImageSize == (int)::fread(pImage, 1, *pImageSize, pMaster));
	} __finally {
		CLOSE_FILE( pMaster );
		/*@hongliang.xin 2009-9-2 */
		::DeleteFile(szPath);
	}

	return pImage;
}

BOOL CMasterImgGen::CalcImageOffset( int NumOfImages,
		                             IMAGE_PARAM * lpImageParam ) 
{
	ASSERT( NumOfImages > 0 );
	ASSERT( NULL != lpImageParam );

	// Boot parameters occupy the first block, skip it.
	int StartSector = m_SctsPerBlock;
	while ( NumOfImages-- > 0 ) {
		FILE * pImage = ::_tfopen( lpImageParam->szPath, _T("rb") );
		if ( NULL == pImage )
			return FALSE;

		// The image data just occupy the main area of a sector.
		int nSectors = ( ::_filelength( _fileno( pImage ) ) + SECTOR_SIZE - 1 ) / SECTOR_SIZE;
		::fclose( pImage );
		pImage = NULL;

		// Every image must start from the boundary of a sector and must
		// occupy space of integral multiple of sectors.
		lpImageParam->Offset = StartSector;   // in sector
		lpImageParam->PaddingSize = nSectors; // in sector

		StartSector += nSectors;
		++lpImageParam;
	}

	return TRUE;
}


BOOL CMasterImgGen::AddBootParam( FILE * pMaster,
		                          int NumOfImages,
					              const IMAGE_PARAM * lpImageParam )
{
	ASSERT( NULL != pMaster );
	ASSERT( NumOfImages > 0 );
	ASSERT( NULL != lpImageParam );

	ASSERT( 0 == ::fflush( pMaster ) );
	ASSERT( 0 == ::_filelength( _fileno( pMaster ) ) );

	BOOL ret = FALSE;
	unsigned char * pBuf = NULL;
	__try {
		int BufSize = sizeof BOOT_PARAM + NumOfImages * sizeof OS_INFO;
		NEW_ARRAY( pBuf, unsigned char, BufSize );

		// Initialize fields of boot parameters.
		BOOT_PARAM * pBoot = ( BOOT_PARAM* )pBuf;
		::memmove( pBoot->Magic, MAGIC, sizeof pBoot->Magic );
		pBoot->Size = SwapWord( sizeof BOOT_PARAM );
		pBoot->TotalSize = SwapWord( ( unsigned short )BufSize );
		::time( &pBoot->TimeStamp );
		pBoot->TimeStamp = SwapDword( pBoot->TimeStamp );
		pBoot->TraceOn = 1;
		pBoot->Reserved = 0;
		pBoot->CurrentOS = ( unsigned char )0;
		pBoot->NumOfOS = ( unsigned char )NumOfImages;
		pBoot->SizeOfOSInfo = SwapWord( sizeof OS_INFO );
		pBoot->OSOffset = SwapWord( sizeof BOOT_PARAM );

		// Initialize OS information.
		OS_INFO * pOS = ( OS_INFO* )( pBuf + sizeof BOOT_PARAM );
		while ( NumOfImages-- > 0 ) {
			::memset( pOS, 0, sizeof *pOS );

			ASSERT( ::strlen( lpImageParam->szName ) < sizeof pOS->Name );
			ASSERT( ::strlen( lpImageParam->szDescription ) < sizeof pOS->Description );
			ASSERT( ::strlen( lpImageParam->szVersion ) < sizeof pOS->Version );
			::strcpy( pOS->Name, lpImageParam->szName );
			::strcpy( pOS->Description, lpImageParam->szDescription );
			::strcpy( pOS->Version, lpImageParam->szVersion );
			// All in sectors.
			pOS->Offset = 
				SwapWord( (unsigned short)lpImageParam->Offset ); 
			pOS->Size = 
				SwapWord( (unsigned short)lpImageParam->PaddingSize );

			++pOS;
			++lpImageParam;
		}
		
		ASSERT(0 == m_SctsPerBlock % 2);
		if ( !WriteBootParam( pMaster, pBuf, BufSize, m_SctsPerBlock / 2 ) )
			__leave;

		// Write backup boot parameter block
		if ( !WriteBootParam( pMaster, pBuf, BufSize, m_SctsPerBlock / 2 ))
			__leave;

		ret = TRUE;
	} __finally {
		DELETE_ARRAY( pBuf );
	}

	return ret;
}
	
BOOL CMasterImgGen::WriteBootParam( FILE * pMaster,
		                            const void * pBootParam,
									int ParamSize,
						            int TotalPages )
{
	// Write boot parameter block to Master Image.
	int nPages = 0;
	const unsigned char * pRemain = (const unsigned char*)pBootParam;
	unsigned char * pMBuf = NULL;
	BOOL ret;

	ret = FALSE;
	__try {
		NEW_ARRAY( pMBuf, unsigned char, SECTOR_SIZE );

	    while ( ParamSize > 0 ) {
			::memset( pMBuf, 0xFF, SECTOR_SIZE );
			int nSize = min( ParamSize, SECTOR_SIZE );			
			::memmove( pMBuf, pRemain, nSize ); //lint !e670

			if ( !WriteMainArea( pMaster, pMBuf ) )
				__leave;

			ParamSize -= nSize;
			pRemain += nSize; //lint !e662
			++nPages;
		}
		ASSERT(nPages <= TotalPages);

		// Padding.
		::memset( pMBuf, 0xFF, SECTOR_SIZE );
		while ( nPages++ < TotalPages ) {
			if ( !WriteMainArea( pMaster, pMBuf ) )
				__leave;
		}
		
		ret = TRUE;
	} __finally {
		DELETE_ARRAY( pMBuf );
	}

	return ret;
}

BOOL CMasterImgGen::MakeImage( FILE * pMaster, 
		                       const IMAGE_PARAM * lpImageParam )
{
	ASSERT( NULL != pMaster );
	ASSERT( NULL != lpImageParam );

	ASSERT( 0 == ::fflush( pMaster ) );
	ASSERT( ::_filelength( _fileno( pMaster ) ) 
		== lpImageParam->Offset * SECTOR_SIZE);

	BOOL ret = FALSE;
	unsigned char * pMBuf = NULL;
	FILE * pImage = NULL;
	__try {
		pImage = ::_tfopen( lpImageParam->szPath, _T("rb") );
		if ( NULL == pImage )
			__leave;
		NEW_ARRAY( pMBuf, unsigned char, SECTOR_SIZE );

		int nPages = 0;
		::memset( pMBuf, 0xFF, SECTOR_SIZE );
		while ( ::fread( pMBuf, sizeof( unsigned char ), SECTOR_SIZE, pImage ) ) {
			if ( !WriteMainArea( pMaster, pMBuf ) )
				__leave;

			++nPages;
			::memset( pMBuf, 0xFF, SECTOR_SIZE );
		}
		ASSERT( 0 == ferror( pImage ) );
		ASSERT( nPages == lpImageParam->PaddingSize );
		
		ret = TRUE;
	} __finally {
		CLOSE_FILE( pImage );
		DELETE_ARRAY( pMBuf );
	}

	return ret;
}


BOOL CMasterImgGen::WriteMainArea( FILE * pMaster,
		                           const unsigned char * pMArea )
{
	ASSERT( NULL != pMaster );
	ASSERT( NULL != pMArea );

	return (SECTOR_SIZE == (int)::fwrite(pMArea, sizeof(unsigned char), SECTOR_SIZE, 
		pMaster));
}

