// RawDiskDlg.cpp : implementation file
//

#include "stdafx.h"
#include "RawDisk.h"
#include "RawDiskDlg.h"
#include "CardIO.h"

#include "debug.h"

#include <winioctl.h>
#include <setupapi.h>
#include <objbase.h>
#include <devguid.h>

#ifdef _DEBUG
#define new DEBUG_NEW
#undef THIS_FILE
static char THIS_FILE[] = __FILE__;
#endif

#define BOOTLOADERSIZE	512
#define WINCESIZE		65024
// program assumes headers are 2 or 7 blocks

#define FLAG_OLDSTYLE		0x00000001
#define FLAG_NEWSTYLE		0x00000002
#define FLAG_HAS_BOOTLOADER	0x00000010
#define FLAG_HAS_WINCE		0x00000020
#define FLAG_HAS_DIAGNOSTIC	0x00000040

#define FLAG_STYLE_MASK		0x0000000F
#define FLAG_HAS_MASK		0x000000F0

/////////////////////////////////////////////////////////////////////////////
// CAboutDlg dialog used for App About

class CAboutDlg : public CDialog
{
public:
	CAboutDlg();

// Dialog Data
	//{{AFX_DATA(CAboutDlg)
	enum { IDD = IDD_ABOUTBOX };
	//}}AFX_DATA

	// ClassWizard generated virtual function overrides
	//{{AFX_VIRTUAL(CAboutDlg)
	protected:
	virtual void DoDataExchange(CDataExchange* pDX);    // DDX/DDV support
	//}}AFX_VIRTUAL

// Implementation
protected:
	//{{AFX_MSG(CAboutDlg)
	//}}AFX_MSG
	DECLARE_MESSAGE_MAP()
};

CAboutDlg::CAboutDlg() : CDialog(CAboutDlg::IDD)
{
	//{{AFX_DATA_INIT(CAboutDlg)
	//}}AFX_DATA_INIT
}

void CAboutDlg::DoDataExchange(CDataExchange* pDX)
{
	CDialog::DoDataExchange(pDX);
	//{{AFX_DATA_MAP(CAboutDlg)
	//}}AFX_DATA_MAP
}

BEGIN_MESSAGE_MAP(CAboutDlg, CDialog)
	//{{AFX_MSG_MAP(CAboutDlg)
	//}}AFX_MSG_MAP
END_MESSAGE_MAP()

/////////////////////////////////////////////////////////////////////////////
// CRawDiskDlg dialog

CRawDiskDlg::CRawDiskDlg(CWnd* pParent /*=NULL*/)
	: CDialog(CRawDiskDlg::IDD, pParent)
{
	//{{AFX_DATA_INIT(CRawDiskDlg)
		// NOTE: the ClassWizard will add member initialization here
	//}}AFX_DATA_INIT
	// Note that LoadIcon does not require a subsequent DestroyIcon in Win32
	m_hIcon = AfxGetApp()->LoadIcon(IDR_MAINFRAME);
}

void CRawDiskDlg::DoDataExchange(CDataExchange* pDX)
{
	CDialog::DoDataExchange(pDX);
	//{{AFX_DATA_MAP(CRawDiskDlg)
	DDX_Control(pDX, IDC_CONTENT, m_cContent);
	DDX_Control(pDX, IDC_SELECT_DISK, m_cSelectDisk);
	//}}AFX_DATA_MAP
}

BEGIN_MESSAGE_MAP(CRawDiskDlg, CDialog)
	//{{AFX_MSG_MAP(CRawDiskDlg)
	ON_WM_SYSCOMMAND()
	ON_WM_PAINT()
	ON_WM_QUERYDRAGICON()
	ON_COMMAND(ID_APP_ABOUT, OnAppAbout)
	ON_CBN_SELCHANGE(IDC_SELECT_DISK, OnSelectDisk)
	ON_BN_CLICKED(IDC_BUTTON_WRITE, OnButtonWrite)
	ON_BN_CLICKED(IDC_BUTTON_READ, OnButtonRead)
	//}}AFX_MSG_MAP
END_MESSAGE_MAP()

/////////////////////////////////////////////////////////////////////////////
// CRawDiskDlg message handlers

UINT PumpMessages()
{
	MSG msg;
	while ( ::PeekMessage( &msg, NULL, 0, 0, PM_NOREMOVE ) ) 
	{ 
		if ( !AfxGetThread()->PumpMessage( ) ) 
		{ 
			AfxPostQuitMessage( 0); 
			return 1;
		}
	}
	return 0;
}

BOOL CRawDiskDlg::OnInitDialog()
{
	CDialog::OnInitDialog();

	// Add "About..." menu item to system menu.

	// IDM_ABOUTBOX must be in the system command range.
	ASSERT((IDM_ABOUTBOX & 0xFFF0) == IDM_ABOUTBOX);
	ASSERT(IDM_ABOUTBOX < 0xF000);

	CMenu* pSysMenu = GetSystemMenu(FALSE);
	if (pSysMenu != NULL)
	{
		CString strAboutMenu;
		strAboutMenu.LoadString(IDS_ABOUTBOX);
		if (!strAboutMenu.IsEmpty())
		{
			pSysMenu->AppendMenu(MF_SEPARATOR);
			pSysMenu->AppendMenu(MF_STRING, IDM_ABOUTBOX, strAboutMenu);
		}
	}

	// Set the icon for this dialog.  The framework does this automatically
	//  when the application's main window is not a dialog
	SetIcon(m_hIcon, TRUE);			// Set big icon
	SetIcon(m_hIcon, FALSE);		// Set small icon

	m_cFont.CreatePointFont(60, "Terminal");
	
	m_cContent.SetFont( &m_cFont );
	GetAllDisksInfo();
	
	return TRUE;  // return TRUE  unless you set the focus to a control
}

void CRawDiskDlg::OnSysCommand(UINT nID, LPARAM lParam)
{
	if ((nID & 0xFFF0) == IDM_ABOUTBOX)
	{
		CAboutDlg dlgAbout;
		dlgAbout.DoModal();
	}
	else
	{
		CDialog::OnSysCommand(nID, lParam);
	}
}

// If you add a minimize button to your dialog, you will need the code below
//  to draw the icon.  For MFC applications using the document/view model,
//  this is automatically done for you by the framework.

void CRawDiskDlg::OnPaint() 
{
	if (IsIconic())
	{
		CPaintDC dc(this); // device context for painting

		SendMessage(WM_ICONERASEBKGND, (WPARAM) dc.GetSafeHdc(), 0);

		// Center icon in client rectangle
		int cxIcon = GetSystemMetrics(SM_CXICON);
		int cyIcon = GetSystemMetrics(SM_CYICON);
		CRect rect;
		GetClientRect(&rect);
		int x = (rect.Width() - cxIcon + 1) / 2;
		int y = (rect.Height() - cyIcon + 1) / 2;

		// Draw the icon
		dc.DrawIcon(x, y, m_hIcon);
	}
	else
	{
		CDialog::OnPaint();
	}
}

// The system calls this to obtain the cursor to display while the user drags
//  the minimized window.
HCURSOR CRawDiskDlg::OnQueryDragIcon()
{
	return (HCURSOR) m_hIcon;
}

void CRawDiskDlg::OnAppAbout() 
{
	CAboutDlg dlgAbout;
	dlgAbout.DoModal();
}

HANDLE OpenDev( const CString& devnam, int mode )
{
	HANDLE hDevice;		 // handle to the drive to be examined 

	CString filename= devnam;	// why doesn't createfile take a 'const' parameter?

	hDevice = CreateFile(filename.GetBuffer(0),		// drive to open
		mode,							// access to the drive
		FILE_SHARE_READ |				// share mode
		FILE_SHARE_WRITE, 
		NULL,							// default security attributes
		OPEN_EXISTING,					// disposition
		0,								// file attributes
		NULL);							// do not copy file attributes

	if (hDevice==INVALID_HANDLE_VALUE) // cannot open the drive
		return NULL;

	return hDevice;
}

bool ReadBlocks( HANDLE hDev, FILE *f, DWORD numblocks, DWORD blocksize, CProgress& cProgress )
{
	bool bResult= true;
	unsigned char *buffer= (unsigned char *)malloc( blocksize );

	DWORD t0= GetTickCount();
	DWORD t1= t0;

	for( DWORD i= 0; i<numblocks && bResult; i++ )
	{
		DWORD readsize;
		bResult= FALSE!=ReadFile( hDev, buffer, blocksize, &readsize, NULL );
		if( bResult && f )
			fwrite( buffer, 1, blocksize, f );

		// update once a second
		if (GetTickCount()/1000!=t1/1000) {
			t1  = GetTickCount();
			cProgress.SetPercentage(100.0*i/numblocks);

			double kbps= 512000.0*i/(t1-t0);
			cProgress.SetSpeed(kbps);
			cProgress.SetRemainingTime((long)(512.0*(numblocks-i)/kbps));
		}
		
		if (cProgress.m_bCancelled)
		{
			bResult= false;
			break;
		}
		PumpMessages();
	}
	return bResult;
}

/*

card with keysignature have the following structure:

block0: starts with string "HTC FLASH KEY\0"
block1: starts with 16 byte md5 sum of hex representation of first 16 bytes of cardid
block2: start with 4 byte value
block3: starts with md5sum of hex representation of first 4 bytes of block2
block4: contains 'N'
( not sure if it really is md5 )

*/

// this is written to block 0 for 5.17
unsigned char magic_flashkey[]=	// "HTC FLASH KEY"
	{0x48,0x54,0x43,0x20,0x46,0x4c,0x41,0x53,0x48,0x20,0x4b,0x45,0x59,0x00,0x00,0x00};

// this is written to block 1
unsigned char magic_thekey[]=
	{0x18,0xdd,0x1f,0x20,0x1d,0x0f,0xd8,0x35,0x98,0xf1,0x72,0x8e,0xa1,0xc8,0x25,0xbf};

unsigned char magic_goldencounter1[]=	/* 0 for now, until we find 1000 */
	{0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00};

unsigned char magic_goldencounter2[]=	/* complement of 0 for now, until we find 1000 */
	{0x9f,0xe5,0xd5,0xef,0x8a,0xe5,0xc5,0x39,0x8e,0xde,0xd8,0x43,0x61,0x1f,0xb1,0x32};

// these are values for block 0 ( or 5 )
unsigned char magic_boot[]=		// "HTC$WALLABY00" - bootloader
	{0x48,0x54,0x43,0x24,0x57,0x41,0x4c,0x4c,0x41,0x42,0x59,0x30,0x30,0x00,0x00,0x00};
unsigned char magic_ce[]=		// "HTC$WALLABY11" - ce
	{0x48,0x54,0x43,0x24,0x57,0x41,0x4c,0x4c,0x41,0x42,0x59,0x31,0x31,0x00,0x00,0x00};
unsigned char magic_ce_boot[]=	// "HTC$WALLABY22" - ce + bootloader
	{0x48,0x54,0x43,0x24,0x57,0x41,0x4c,0x4c,0x41,0x42,0x59,0x32,0x32,0x00,0x00,0x00};
unsigned char magic_diagnostic[]=// "HTC$WALLABY33" - diagnostic card
	{0x48,0x54,0x43,0x24,0x57,0x41,0x4c,0x4c,0x41,0x42,0x59,0x33,0x33,0x00,0x00,0x00};
// "HTC$WALLABY44" - gsm
// "HTC$WALLABY55" - gsm + ce
unsigned char magic_unknown[]=			/* block 4 is unknown but needs 4E */
	{0x4E,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00};

BOOL TestHeader( unsigned char *buffer, DWORD blocksize, DWORD &blocks, DWORD &flags )
{
	blocks= 0;
	flags= 0;

	if( !memcmp( buffer, magic_flashkey, 16 ) )
	{
		// certified new style header
		blocks= 7;
		buffer+= 5*blocksize;
		flags|= FLAG_NEWSTYLE;
	}
	else
	{
		// assume old style header
		blocks= 2;
		flags|= FLAG_OLDSTYLE;
	}

	if( !memcmp( buffer, magic_boot, 16 ) )
	{
		flags|= FLAG_HAS_BOOTLOADER;
	}
	else if( !memcmp( buffer, magic_ce, 16 ) )
	{
		flags|= FLAG_HAS_WINCE;
	}
	else if( !memcmp( buffer, magic_ce_boot, 16 ) )
	{
		flags|= FLAG_HAS_WINCE|FLAG_HAS_BOOTLOADER;
	}
	else if( !memcmp( buffer, magic_diagnostic, 16 ) )
	{
		flags|= FLAG_HAS_DIAGNOSTIC;
		blocks--;
	}

	if( !(flags&FLAG_HAS_MASK))
	{
		AfxMessageBox( "Did not find correct flash signature, aborting!", MB_ICONHAND );
		return false;
	}

	return true;
}

bool WriteBlock( HANDLE hDev, unsigned char *buf, DWORD blocksize, DWORD numblocks )
{
	bool bResult= true;

	for( DWORD i= 0; i<numblocks && bResult; i++ )
	{
		DWORD writesize;
		if (!WriteFile( hDev, buf, blocksize, &writesize, NULL ))
		{
			error("WriteFile");
			bResult= false;
		}
		buf+= blocksize;
	}
	return bResult;
}

bool WriteBlocksFromFile( HANDLE hDev, DWORD blocksize, DWORD numblocks, DWORD *writtenblocks, const char *filename, int part, CProgress& cProgress )
{
	FILE *f= fopen( filename, "rb" );
	*writtenblocks= 0;
	if( !f )
	{
		error("fopen %s", filename);
		return false;
	}

	if( stricmp( filename+strlen(filename)-4, ".nbf" )==0 )
	{
		if( part==1 )
		{
			// skip bootloader part + 32 bytes
			fseek( f, (BOOTLOADERSIZE*blocksize)+32, SEEK_SET );
		}
		else
		{
			// skip first 32 bytes
			fseek( f, 32, SEEK_SET );
		}
	}
	else if( stricmp( filename+strlen(filename)-4, ".img" )==0 )
	{
		// combined file with raw header - interpret header to find our data
		BYTE *buffer= new BYTE[7 * blocksize];
		fread( buffer, blocksize, 7, f );
		fseek( f, 0, SEEK_SET );

		DWORD blocks;
		DWORD flags;

		if( !TestHeader( buffer, blocksize, blocks, flags ) )
		{
			delete buffer;
			fclose( f );
			return false;
		}
		delete buffer;

		if( part==0 )
		{
			if( !(flags&FLAG_HAS_BOOTLOADER) )
			{
				fclose( f );
				CString s;
				s.Format( "File %s does not contain bootloader!", filename );
				AfxMessageBox( s, MB_ICONHAND );
				return false;
			}
		}
		else if( part==1 )
		{
			if( !(flags&FLAG_HAS_WINCE) )
			{
				fclose( f );
				CString s;
				s.Format( "File %s does not contain Windows CE!", filename );
				AfxMessageBox( s, MB_ICONHAND );
				return false;
			}
			if( flags&FLAG_HAS_BOOTLOADER )
				blocks+= BOOTLOADERSIZE;
		}
		fseek( f, blocks*blocksize, SEEK_SET );
	}

	BYTE *buffer= new BYTE[blocksize];

	DWORD t0  = GetTickCount();
	DWORD t1  = t0;

	bool bResult= true;

	*writtenblocks= 0;
	for( DWORD i= 0; (numblocks==0 || i<numblocks) && bResult; i++ )
	{
		memset( buffer, 0, blocksize );

		DWORD nBlocksread= fread( buffer, blocksize, 1, f );
		if( numblocks==0 && nBlocksread==0 )
			break;
		bResult= WriteBlock( hDev, buffer, blocksize, 1 );
		(*writtenblocks)++;
		
		// update once a second
		if (GetTickCount()/1000!=t1/1000) {
			t1  = GetTickCount();
			cProgress.SetPercentage(100.0*i/numblocks);

			double kbps= 512000.0*i/(t1-t0);
			cProgress.SetSpeed(kbps);
			cProgress.SetRemainingTime((long)(512.0*(numblocks-i)/kbps));
		}
		if (cProgress.m_bCancelled)
		{
			bResult= false;
			break;
		}

		PumpMessages();
	}

	delete buffer;
	fclose( f );
	return bResult;
}

void CloseDev( HANDLE hDev )
{
	if( hDev!=NULL )
		CloseHandle( hDev );	
}

BOOL GetDriveGeometry(const CString& devicename, DISK_GEOMETRY *pdg)
{
	HANDLE hDevice;		 // handle to the drive to be examined 
	BOOL bResult;		 // results flag
	DWORD junk;			 // discard results

	hDevice = OpenDev( devicename, 0 );

	if (hDevice == INVALID_HANDLE_VALUE) // cannot open the drive
		return (FALSE);

	bResult = DeviceIoControl(hDevice,	// device to be queried
		IOCTL_DISK_GET_DRIVE_GEOMETRY,	// operation to perform
		NULL, 0, // no input buffer
		pdg, sizeof(*pdg),		 // output buffer
		&junk,								 // # bytes returned
		(LPOVERLAPPED) NULL);	// synchronous I/O

	CloseDev( hDevice );

	return bResult;
}

BOOL ReadRaw(const CString& devicename, DWORD size, unsigned char * buffer, DWORD bufsize )
{
	HANDLE hDevice;		 // handle to the drive to be examined 
	BOOL bResult;		 // results flag
	HGLOBAL hMem;
	char *pMem;
	DWORD readsize;

	hDevice = OpenDev(devicename, GENERIC_READ );

	if (hDevice==INVALID_HANDLE_VALUE) // cannot open the drive
		return false;

	hMem= GlobalAlloc( GHND, size );
	if( hMem==NULL )
		return false;
	pMem= (char*)GlobalLock(hMem);

	bResult= FALSE!=ReadFile( hDevice, pMem, size, &readsize, NULL );

	CloseDev(hDevice);

	memcpy( buffer, pMem, bufsize );
	GlobalUnlock(hMem);

	GlobalFree( hMem );
	return bResult;
}

BOOL ReadDiskInfo( CDiskInfo *dip )
{

	DISK_GEOMETRY pdg;						// disk drive geometry structure
	BOOL bResult;								 // generic results flag

	bResult = GetDriveGeometry( dip->devicename, &pdg );
	if( bResult )
	{
		dip->sectorsize= pdg.BytesPerSector;
		dip->mediatype= pdg.MediaType;

		bResult= ReadRaw( dip->devicename, pdg.BytesPerSector, dip->firstdata, sizeof( dip->firstdata ) );
	}

	return bResult;
}

#define BLOCKSIZE 16

CString DumpMem( unsigned char *buffer, int bufsize )
{
	CString res, ha;
	int i,j;
	res= "";

	for( i= 0; i<bufsize; i+= BLOCKSIZE )
	{
		for( j= 0; j<min(BLOCKSIZE, bufsize-i); j++ )
		{
			ha.Format( "%02X ", buffer[i+j] );
			res+= ha;
		}
		for( j= 0; j<min(BLOCKSIZE, bufsize-i); j++ )
		{
			if( isprint(buffer[i+j]) )
			{
				ha.Format( "%c", buffer[i+j] );
				res+= ha;
			}
			else
			{
				ha.Format( "." );
				res+= ha;
			}
		}
		res+= "\r\n";
	}

	return res;
}


void CRawDiskDlg::OnSelectDisk() 
{
	int i;
	i= m_cSelectDisk.GetCurSel();
	if( i>=0 )
	{
		CDiskInfo *dip= (CDiskInfo *)m_aDiskInfo.GetAt( m_cSelectDisk.GetItemData(i) );
		if( ReadDiskInfo( dip ) )
		{
			m_cContent.SetWindowText( DumpMem(dip->firstdata, sizeof(dip->firstdata)) );
		}
		else 
		{
			m_cContent.SetWindowText("");
		}
	}
}

// !!!!!!!!!!!!!!!!!!!!!!!!!!!
// BIG BUG WARNING, both 'index' and 'Id' have no relation to the actual disk id.
// this code should enumerate the correct registry keys instead of guessing
//
HDEVINFO DoDeviceEnum( GUID InterfaceClassGuid )
{
    HDEVINFO DeviceInfoSet;
    HDEVINFO NewDeviceInfoSet;
	DWORD Err;

    // Create a device information set that will be the container for 
    // the device interfaces.

    DeviceInfoSet = SetupDiCreateDeviceInfoList(NULL, NULL);
    
    if(DeviceInfoSet == INVALID_HANDLE_VALUE) {
        Err = GetLastError();
        printf( "SetupDiCreateDeviceInfoList failed: %lx.\n", Err );
        return 0;
    }

    // Retrieve the device information set for the interface class.

    NewDeviceInfoSet = SetupDiGetClassDevsEx( &InterfaceClassGuid,
        NULL,
        NULL,
		DIGCF_PRESENT,
//		DIGCF_PRESENT | DIGCF_DEVICEINTERFACE,
//		DIGCF_PRESENT | DIGCF_ALLCLASSES,
        DeviceInfoSet,
        NULL,
        NULL
    );

    if(NewDeviceInfoSet == INVALID_HANDLE_VALUE) 
    {
        Err = GetLastError();
        printf( "SetupDiGetClassDevsEx failed: %lx.\n", Err );
        return 0;
    }

    return NewDeviceInfoSet;
}
CString GetDevInfoString(HDEVINFO hDevInfo, SP_DEVINFO_DATA *pDeviceInfoData, DWORD dwProperty)
{
	DWORD dwDataType;
	BYTE *buffer = NULL;
	DWORD buffersize = 0;

	while (!SetupDiGetDeviceRegistryProperty(hDevInfo,
		pDeviceInfoData, dwProperty, &dwDataType,
		buffer,	buffersize,	&buffersize))
	{
		if (GetLastError() == ERROR_INSUFFICIENT_BUFFER)
		{
			// Change the buffer size.
			if (buffer)
				LocalFree(buffer);
			buffer = (BYTE*)LocalAlloc(LPTR,buffersize);
		}
		else
		{
			// Insert error handling here.
			break;
		}
	}

	CString result;
	if (buffer)
	{
		if (dwDataType==REG_SZ)
			result= buffer;
		else if (dwDataType==REG_DWORD)
			result.Format("DWORD:%08lx", *(DWORD*)buffer);
		else if (dwDataType==REG_MULTI_SZ)
		{
			CString tmp;
			for (int i=0 ; i<buffersize ; i++)
			{
				tmp.Format(" '%s'", buffer+i);
				result += tmp;

				i+=strlen((char*)buffer);
			}
		}
		else {
			result.Format("regtype %08lx", dwDataType);
			CString tmp;
			for (int i=0 ; i<buffersize ; i++)
			{
				tmp.Format(" %02x", buffer[i]);
				result += tmp;
			}
		}

		LocalFree(buffer);
		buffer= NULL;
	}
	buffersize = 0;

	return result;
}

/*
todo: much easier way to find disks ( and actually works )

- look in HKLM\System\CurrentControlSet\Services\Disk\Enum
    keys  0..5  contain reference to :

refer to HKLM\System\CurrentControlSet\Enum\[reference]\<numerickey>
               "FriendlyName" -> ....

not sure about the portability to non-nt platforms.

*/
void CRawDiskDlg::GetAllDisksInfo()
{
	HDEVINFO hDevInfo= DoDeviceEnum( GUID_DEVCLASS_DISKDRIVE );
	SP_DEVINFO_DATA deviceInfoData;
	int i;
	CDiskInfo *dip;

	// Enumerate through all devices in Set.
	deviceInfoData.cbSize = sizeof(SP_DEVINFO_DATA);
	for( i=0; SetupDiEnumDeviceInfo(hDevInfo,i,&deviceInfoData); i++ )
	{
		dip= new CDiskInfo;

		dip->Driver= GetDevInfoString(hDevInfo, &deviceInfoData, SPDRP_DRIVER);
		dip->FriendlyName= GetDevInfoString(hDevInfo, &deviceInfoData, SPDRP_FRIENDLYNAME);

		m_aDiskInfo.Add( dip );
	}

	if ( GetLastError()!=NO_ERROR &&
		GetLastError()!=ERROR_NO_MORE_ITEMS )
	{
		// Insert error handling here.
	}

	int j;
	for( i= 0; i<m_aDiskInfo.GetSize(); i++ )
	{
		dip= (CDiskInfo *)m_aDiskInfo.GetAt( i );
		j= dip->Driver.ReverseFind('\\');
		if( j>=0 )
		{
			CString ha= dip->Driver.Mid( j+1 );
			dip->Id= atoi( ha );
			dip->index= i;
		}
		else
		{
			dip->Id= -1;
			dip->index= i;
		}

		dip->devicename.Format("\\\\.\\PhysicalDrive%d", dip->index );
	}

	for( i= 0; i<m_aDiskInfo.GetSize(); i++ )
	{
		dip= (CDiskInfo *)m_aDiskInfo.GetAt( i );

		if( dip->Id>=0 )
		{
			CString s;
			s.Format( "disk %u - %s", dip->Id, (const char*)(dip->FriendlyName) );
			m_cSelectDisk.SetItemData( m_cSelectDisk.AddString( s ), i );
		}
	}
	m_cSelectDisk.SetCurSel( m_cSelectDisk.GetCount()-1 );
	if( m_cSelectDisk.GetCount()<=0 )
		AfxMessageBox( "Could not find any disk drives!" );

	OnSelectDisk();
}

BOOL ReadImage( CDiskInfo *dip, DWORD blocksize, char *Bootloader, char *WinCE )
{
	unsigned char *buffer;
	buffer= (unsigned char *)malloc( blocksize*7 );
	BOOL bResult;
	DWORD blocks= 0, HeaderBlocks= 0;
	DWORD BootloaderBlocks= 0, WinCEBlocks= 0;
	DWORD BootloaderOffset= 0, WinCEOffset= 0;

	CProgress cProgress;
	cProgress.Create();

	bResult= ReadRaw( dip->devicename, blocksize*7, buffer, blocksize*7 );
	if( !bResult )
	{
		CString s;
		s.Format( "Failed to open and read device %s", dip->devicename );

		cProgress.DestroyWindow();

		AfxMessageBox( s, MB_ICONHAND );
		return bResult;
	}

	DWORD flags;
	bResult= TestHeader( buffer, blocksize, HeaderBlocks, flags );
	if( !bResult )
	{
		cProgress.DestroyWindow();
		return bResult;
	}
	blocks= HeaderBlocks;

	if( flags&FLAG_HAS_BOOTLOADER )
	{
		BootloaderOffset= blocks;
		BootloaderBlocks= BOOTLOADERSIZE;
		blocks+= BOOTLOADERSIZE;
	}
	if( flags&FLAG_HAS_WINCE )
	{
		WinCEOffset= blocks;
		WinCEBlocks= WINCESIZE;
		blocks+= WINCESIZE;
	}

	FILE *f1= NULL, *f2= NULL;

	if( Bootloader && *Bootloader && stricmp( Bootloader+strlen(Bootloader)-4, ".img" )==0 )
	{
		// dump to one file; make it the read image format
		// no support for nbf files, for now, just raw img files
		f1= fopen( Bootloader, "wb" );
		f2= f1;
		fwrite( buffer, blocksize, HeaderBlocks, f1 );
	}
	else if( WinCE && *WinCE && stricmp( WinCE+strlen(WinCE)-4, ".img" )==0 )
	{
		// dump to one file; make it the read image format
		// no support for nbf files, for now, just raw img files
		f1= fopen( WinCE, "wb" );
		f2= f1;
		fwrite( buffer, blocksize, HeaderBlocks, f1 );
	}
	else
	{
		if( *Bootloader && BootloaderBlocks )
			f1= fopen( Bootloader, "wb" );
		if( *WinCE && WinCEBlocks )
			f2= fopen( WinCE, "wb" );
	}

	free( buffer );

	HANDLE hDev;
	hDev= OpenDev( dip->devicename, GENERIC_READ );

	if( hDev==NULL )
	{
		cProgress.DestroyWindow();

		AfxMessageBox( "Failed to open device " + dip->devicename, MB_ICONHAND );
		return false;
	}

	bool res= ReadBlocks( hDev, NULL, HeaderBlocks, blocksize, cProgress ); 
	if( res && BootloaderBlocks )
	{
		res= ReadBlocks( hDev, f1, BootloaderBlocks, blocksize, cProgress );
	}
	if( res && WinCEBlocks )
	{
		res= ReadBlocks( hDev, f2, WinCEBlocks, blocksize, cProgress );
	}
	CloseDev( hDev );
	if( f1 )
		fclose( f1 );
	if( f2 && f1!=f2 )
		fclose( f2 );

	cProgress.DestroyWindow();

	if (res) {
		CString s;
		s.Format( "Successfully copied image of %lu blocks", blocks );
		AfxMessageBox( s, MB_ICONHAND );
	}
	else {
		CString s;
		s.Format( "cancelled" );
		AfxMessageBox( s, MB_ICONHAND );
	}

	return true;
}

void ReallyCalcCrc( DWORD *crc, unsigned char *buf, DWORD size )
{
	DWORD i;
	for( i= 0; i<size; i+= 4 )
		*crc+= *(DWORD*)( buf+i );
}

BOOL CalcCrc( DWORD *crc, DWORD blocksize, DWORD numblocks, const char *filename, int part )
{
	FILE *f= fopen( filename, "rb" );
	BOOL bResult;
	if( !f )
		return false;

	unsigned char *buffer;

	if( stricmp( filename+strlen(filename)-4, ".nbf" )==0 )
	{
		if( part==1 )
		{
			// skip bootloader part + 32 bytes
			fseek( f, (BOOTLOADERSIZE*blocksize)+32, SEEK_SET );
		}
		else
		{
			// skip first 32 bytes
			fseek( f, 32, SEEK_SET );
		}
	}
	else if( stricmp( filename+strlen(filename)-4, ".img" )==0 )
	{
		buffer= (unsigned char *)malloc( 7 * blocksize );
		fread( buffer, blocksize, 7, f );
		fseek( f, 0, SEEK_SET );

		DWORD blocks;
		DWORD flags;
		bResult= TestHeader( buffer, blocksize, blocks, flags );
		if( !bResult )
		{
			free( buffer );
			fclose( f );
			return bResult;
		}
		free( buffer );

		if( part==1 )
		{
			if( flags&FLAG_HAS_BOOTLOADER )
				blocks+= BOOTLOADERSIZE;
		}

		fseek( f, blocks*blocksize, SEEK_SET );
	}

	buffer= (unsigned char *)malloc( blocksize );
	DWORD i;

	for( i= 0; i<numblocks; i++ )
	{
		memset( buffer, 0, blocksize );
		fread( buffer, 1, blocksize, f );
		ReallyCalcCrc( crc, buffer, blocksize );

		PumpMessages();
	}

	fclose( f );
	free( buffer );
	return true;
}

BOOL WriteImage( CDiskInfo* dip, DWORD blocksize, UINT headerstyle, const char *Bootloader, const char *WinCE, const char *Diag )
{
	unsigned char *buffer= (unsigned char *)malloc( blocksize );

	DWORD crc= 0;
	DWORD blocks= 0;

	CProgress cProgress;
	cProgress.Create();

	if( Bootloader && *Bootloader )
	{
		if( !CalcCrc( &crc, blocksize, BOOTLOADERSIZE, Bootloader, 0 ) )
		{
			CString s;
			s.Format( "Failed to calculate CRC from %s", Bootloader );

			cProgress.DestroyWindow();

			AfxMessageBox( s, MB_ICONHAND );
			return false;
		}
	}

	if( WinCE && *WinCE )
	{
		if( !CalcCrc( &crc, blocksize, WINCESIZE, WinCE, 1 ) )
		{
			CString s;
			s.Format( "Failed to calculate CRC from %s", Bootloader );

			cProgress.DestroyWindow();

			AfxMessageBox( s, MB_ICONHAND );
			return false;
		}
	}

	HANDLE hDev= OpenDev( dip->devicename, GENERIC_WRITE );
	if( hDev==NULL )
	{
		cProgress.DestroyWindow();

		AfxMessageBox( "Failed to open device " + dip->devicename, MB_ICONHAND );

		return false;
	}

	if( headerstyle )
	{
		// long header
		memset( buffer, 0, blocksize );
		memcpy( buffer, magic_flashkey, 16 );
		WriteBlock( hDev, buffer, blocksize, 1 );   // block 0

		memset( buffer, 0, blocksize );
		memcpy( buffer, magic_thekey, 16 );
		WriteBlock( hDev, buffer, blocksize, 1 );   // block 1
		
		memset( buffer, 0, blocksize );
		memcpy( buffer, magic_goldencounter1, 16 );
		WriteBlock( hDev, buffer, blocksize, 1 );   // block 2
		
		memset( buffer, 0, blocksize );
		memcpy( buffer, magic_goldencounter2, 16 );
		WriteBlock( hDev, buffer, blocksize, 1 );   // block 3
		
		memset( buffer, 0, blocksize );
		memcpy( buffer, magic_unknown, 16 );
		WriteBlock( hDev, buffer, blocksize, 1 );   // block 4
		
		blocks+= 5;
	}

	memset( buffer, 0, blocksize );
	if( Diag && *Diag )
		memcpy( buffer, magic_diagnostic, 16 );
	else if( WinCE && *WinCE && Bootloader && *Bootloader )
		memcpy( buffer, magic_ce_boot, 16 );
	else if( WinCE && *WinCE )
		memcpy( buffer, magic_ce, 16 );
	else if( Bootloader && *Bootloader )
		memcpy( buffer, magic_boot, 16 );
	WriteBlock( hDev, buffer, blocksize, 1 );
	blocks++;

	if( !Diag || !*Diag )
	{
		memset( buffer, 0, blocksize );
		memcpy( buffer, &crc, 4 );
		WriteBlock( hDev, buffer, blocksize, 1 );
		blocks++;
	}

	free( buffer );

	DWORD actualblocks=0;
	bool bResult= true;
	if( bResult && Bootloader && *Bootloader )
	{
		bResult= WriteBlocksFromFile( hDev, blocksize, BOOTLOADERSIZE, &actualblocks, Bootloader, 0, cProgress );
		blocks+= actualblocks;
	}

	if( bResult && WinCE && *WinCE )
	{
		bResult= WriteBlocksFromFile( hDev, blocksize, WINCESIZE, &actualblocks, WinCE, 1, cProgress );
		blocks+= actualblocks;
	}

	if( bResult && Diag && *Diag )
	{
		DWORD newblocks= WriteBlocksFromFile( hDev, blocksize, 0, &actualblocks, Diag, 2, cProgress );
		blocks+= actualblocks;
	}

	CloseDev( hDev );

	cProgress.DestroyWindow();

	if (bResult) {
		CString s;
		s.Format( "Successfully copied image of %lu blocks", blocks );
		AfxMessageBox( s, MB_ICONHAND );
	}
	else {
		CString s;
		s.Format( "cancelled", blocks );
		AfxMessageBox( s, MB_ICONHAND );
	}

	return true;
}

void CRawDiskDlg::OnButtonRead() 
{
	int i;
	BOOL bResult;
	CDiskInfo *dip;

	i= m_cSelectDisk.GetCurSel();
	i= m_cSelectDisk.GetItemData( i );

	CReadCard cw;
	cw.m_bBootloader= true;
	cw.m_bWinCE= true;
	if( cw.DoModal()==IDOK )
	{
		dip= (CDiskInfo *)m_aDiskInfo.GetAt( i );
		// just do it!
		bResult= ReadImage( dip, dip->sectorsize, cw.m_bBootloader?(char*)(const char*)cw.m_sBootloader:"", cw.m_bWinCE?(char*)(const char*)cw.m_sWinCE:"" );
	}
}

void CRawDiskDlg::OnButtonWrite() 
{
	int i;
	BOOL bResult;
	CDiskInfo *dip;

	i= m_cSelectDisk.GetCurSel();
	i= m_cSelectDisk.GetItemData( i );

	dip= (CDiskInfo *)m_aDiskInfo.GetAt( i );
	if( !dip->Id )
	{
		if( AfxMessageBox( "This will probably overwrite your boot drive. Are you SURE?", MB_OKCANCEL|MB_ICONINFORMATION|MB_DEFBUTTON2 )!=IDOK )
			return;
	}
	else if (dip->mediatype!=RemovableMedia) {
		if( AfxMessageBox( "This is not a removable disk, this will probably overwrite a harddisk. Are you SURE?", MB_OKCANCEL|MB_ICONINFORMATION|MB_DEFBUTTON2 )!=IDOK )
			return;
	}

	CWriteCard cw;
	cw.m_bHeader= true;
	cw.m_bWinCE= true;
	cw.m_iHeaderType= 0;
	if( cw.DoModal()==IDOK )
	{
		if( !dip->Id )
		{
			if( AfxMessageBox( "Are you REALLY SURE you want to overwrite your boot drive?", MB_OKCANCEL|MB_ICONSTOP|MB_DEFBUTTON2 )!=IDOK )
				return;
		}
		// just do it!
		bResult= WriteImage( dip, dip->sectorsize, cw.m_iHeaderType,
			cw.m_bBootloader?cw.m_sBootloader:"",
			cw.m_bWinCE?cw.m_sWinCE:"",
			cw.m_bDiag?cw.m_sDiag:"" );

		OnSelectDisk();
	}
}