꿈/희망/절망

// 2013. 08. 02
// mangg@manggong.org
// 
// 
// 가상이미지 파일에 부트로더 복사.
// 할일 없어서 이거나 주물럭 주물럭.

#include "stdafx.h"
#include <Windows.h>

#define PARTITION_POS	0x1BE

#pragma pack(push, 1)
typedef struct _PARTITION                                                                          
{                                                                                                  
	BYTE ACTIVE;	// 부팅 가능 플래그 (0x80 : 부팅 가능, 0x00 : 부팅 불가)
	BYTE CHS_BEGIN[3];
	BYTE TYPE;		// 파티션 타입 ( 0x07:NTFS, 0x05:Extended Partition, 0x0B:FAT32, 0x0E:FAT16, 0x0F:Extended Partition, exp INT 13h
					//				0x1B: Hidden FAT32, 0x42:Secure FileSystem, 0x82:Linux Swap partition, 0x83:Linux native file systems)
	BYTE CHS_END[3];
	INT32 LBA_START;
	INT32 LENGTH;
} PARTITION, *PPARTITION;

typedef struct _FAT
{
	BYTE CODE[3];
	CHAR SYSTEMID[8];
	WORD SECTOR_SIZE;
	BYTE SECTOR_CLUSTER;
	WORD RESERVED_SECTOR;
	BYTE FAT_NUM;
	WORD ROOT_DIRECTORY_ENTRY_CNT;
	WORD TOTAL_SECTOR16;
	BYTE MEDIATYPE;
	WORD SECTOR_FAT16;
	WORD SECTOR_TRACK;
	WORD SECTOR_HEAD;
	DWORD HIDDEN_SECTOR;
	DWORD TOTAL_SECTOR32;
	DWORD SECTOR_FAT32;
	WORD EXT_FLAGS;
	WORD VER;
	DWORD ROOT_DIRECTORY_CLUSTER;
	WORD INFO;
	WORD BACKUP_BOOT_SECTOR;
	BYTE RESERVED_ID[12];
	BYTE PHYSICALDISK_TYPE;
	BYTE RESERVED;
	BYTE SIGNATURE;
	DWORD VOLUMEID;
	BYTE VOLUME_LABLE[11];
	BYTE FILESYSTEM_TYPE[8];
} FAT32, *PFAT32;


#pragma pack(pop)

bool LoadBootloader(BYTE *pDiskInfo,TCHAR *BootloaderFile);
bool CopyBootloader(HANDLE hFile, __int64 sector, BYTE *pDiskInfo,TCHAR *BootloaderFile);
bool GetFAT(HANDLE hFile,  __int64 sector, BYTE *pDiskInfo, PFAT32 FAT_INFO);
bool GetActivePartition(HANDLE hFile, BYTE *pDiskInfo, PPARTITION PartitionInfo);
bool DISK_WRITE(HANDLE hFile, __int64 sector, BYTE *pData, int sectorSize);
bool DISK_WRITE_PARTICLE(HANDLE hFile, __int64 sector, DWORD offset, BYTE *pData, int writeSize);
bool DISK_READ(HANDLE hFile, __int64 sector, BYTE *pData, int sectorSize);

void Print_PhysicalList()
{
	TCHAR physical[65536];
	TCHAR logical[65536];

    QueryDosDevice(NULL, physical, sizeof(physical));
	printf("Physical Drive Number : \n");

    for (char *pos = physical; *pos; pos+=strlen(pos)+1) 
	{
        QueryDosDevice(pos, logical, sizeof(logical));

		if( strstr(pos, "Physical") != NULL )
		{
			printf("%s : %s\n", pos, logical);
		}
    }
}

int _tmain(int argc, _TCHAR* argv[])
{
	HANDLE hFile;
	TCHAR szDrive[MAX_PATH];
	DWORD	FAT_SECTOR = 0;
	BYTE	*pDiskInfo;
	PARTITION PartitionInfo;
	FAT32 FAT_INFO;

	if( argc < 3 )
	{
		printf("Usage : bootcopy bootloader VirtualFile\n\n");
		//Print_PhysicalList();
		return -1;
	}

	_stprintf(szDrive, _T("%s"), argv[2]);

	hFile = ::CreateFile(
						szDrive, 
						GENERIC_READ | GENERIC_WRITE, 
						FILE_SHARE_READ | FILE_SHARE_WRITE, 
						NULL, 
						OPEN_EXISTING, 
						0, 
						NULL);

	if( hFile == INVALID_HANDLE_VALUE )
	{
		printf("Don't read drive\n");
		return -1;
	}


	pDiskInfo = (BYTE*)malloc(0x200);
	if(pDiskInfo == NULL )
	{
		CloseHandle(hFile);
		printf("Don't Allocate Memory\n");
		return -1;
	}

	DWORD junk;
	BOOL bResult;


	char bs[255];
	memset(bs,0,255);

	/////////////////////////////////////////
	// READ PARTITION INFO

	if( GetActivePartition(hFile, pDiskInfo, &PartitionInfo) == false )
	{
		CloseHandle(hFile);
		free(pDiskInfo);
		printf("Don't Read Partition Info\n");
		return -1;
	}

	/////////////////////////////////////////
	// READ FAT32 INFO
	if( GetFAT(hFile, PartitionInfo.LBA_START, pDiskInfo, &FAT_INFO) == false )
	{
		CloseHandle(hFile);
		free(pDiskInfo);
		printf("Don't Read FAT Info\n");
		return -1;
	}

	/////////////////////////////////////////
	// COPY BOOTLOADER
	if( CopyBootloader(hFile, PartitionInfo.LBA_START, pDiskInfo, argv[1]) == false )
	{
		CloseHandle(hFile);
		free(pDiskInfo);
		printf("Don't Copy Bootloader\n");
		return -1;
	}

	printf("Copy Success...\n");

	free(pDiskInfo);
	CloseHandle(hFile);

	return 0;
}

bool DISK_READ(HANDLE hFile, __int64 sector, BYTE *pData, int sectorSize)
{
	DWORD read = 0;

	LARGE_INTEGER li;

	li.QuadPart = (0x200 * sector) & ~(0x200-1);

	if (::SetFilePointer(hFile, li.LowPart, &li.HighPart, FILE_BEGIN) == INVALID_SET_FILE_POINTER)
	{
		::GetLastError();
		return false;
	}

	if (!::ReadFile(hFile, pData, sectorSize, &read, NULL))
		return false;

	return true;
}

bool DISK_WRITE(HANDLE hFile, __int64 sector, BYTE *pData, int sectorSize)
{
	DWORD written = 0;

	LARGE_INTEGER li;

	li.QuadPart = (0x200 * sector) & ~(0x200-1);

	if (::SetFilePointer(hFile, li.LowPart, &li.HighPart, FILE_BEGIN) == INVALID_SET_FILE_POINTER)
	{
		::GetLastError();
		return false;
	}

	if (!::WriteFile(hFile, pData, sectorSize, &written, NULL))
		return false;

	return true;
}

bool DISK_WRITE_PARTICLE(HANDLE hFile, __int64 sector, DWORD offset, BYTE *pData, int writeSize)
{
	DWORD written = 0;

	LARGE_INTEGER li;

	li.QuadPart = (0x200 * sector) & ~(0x200-1);
	li.LowPart += offset;

	if (::SetFilePointer(hFile, li.LowPart, &li.HighPart, FILE_BEGIN) == INVALID_SET_FILE_POINTER)
	{
		printf("SetFilePointer error : %d\n",::GetLastError());
		return false;
	}

	if (!::WriteFile(hFile, pData, writeSize, &written, NULL))
	{
		return false;
	}

	return true;
}

bool GetFAT(HANDLE hFile,  __int64 sector, BYTE *pDiskInfo, PFAT32 FAT_INFO)
{
	DWORD read = 0;

	if(DISK_READ(hFile, sector, pDiskInfo, 0x200) == false )
	{
		return false;
	}

	memcpy(FAT_INFO, pDiskInfo, sizeof(FAT32));

	return true;
}

bool GetActivePartition(HANDLE hFile, BYTE *pDiskInfo, PPARTITION PartitionInfo)
{
	PPARTITION PartitionTemp = NULL;

	if( DISK_READ(hFile, 0, pDiskInfo, 0x200) == false )
		return false;

	for(int i = 0; i < 4; i++ )
	{
		PartitionTemp = (PPARTITION) (pDiskInfo + PARTITION_POS + (sizeof(PARTITION) * i));

		if(PartitionTemp->ACTIVE == 0x80 )
		{
			memcpy(PartitionInfo, PartitionTemp, sizeof(PARTITION));
			return true;
		}
	}

	return false;
}

bool CopyBootloader(HANDLE hFile, __int64 sector, BYTE *pDiskInfo,TCHAR *BootloaderFile)
{
	if( LoadBootloader(pDiskInfo, BootloaderFile) == false)
		return false;

	if( DISK_WRITE_PARTICLE(hFile, sector, 0, pDiskInfo, 3) == false )
	{
		printf("ERROR : Copy \"BOOT JMP CODE\" to DISK\n");
		return false;
	}

	pDiskInfo += 0x5A;
	if( DISK_WRITE_PARTICLE(hFile, sector, 0x5A, pDiskInfo, 0x200-0x5A) == false)
	{
		printf("ERROR : Copy \"BOOT CODE\" to DISK\n");
		return false;
	}

	return true;
}

bool LoadBootloader(BYTE *pDiskInfo,TCHAR *BootloaderFile)
{
	HANDLE hBOOTFile;
	DWORD dwReadn;

	hBOOTFile = ::CreateFile(
						BootloaderFile, 
						GENERIC_READ | GENERIC_WRITE, 
						FILE_SHARE_READ | FILE_SHARE_WRITE, 
						NULL, 
						OPEN_EXISTING, 
						0, 
						NULL);

	if( hBOOTFile == INVALID_HANDLE_VALUE )
	{
		printf("Don't open bootloader\n");
		return false;
	}

	if(!::ReadFile(hBOOTFile, pDiskInfo, 0x200, &dwReadn, NULL))
	{
		printf("Don't read bootloader\n");

		CloseHandle(hBOOTFile);
		return false;
	}
	if(dwReadn != 0x200)
	{
		printf("Is not match bootloader size\n");
		CloseHandle(hBOOTFile);
		return false;
	}

	CloseHandle(hBOOTFile);
	return true;
}