#include "stdafx.h"

#include "debug.h"
#include "diskfunctions.h"

#include "WINDEV.H"
#include "diskio.h"

PhysicalDisk::PhysicalDisk()
{
	m_sid= NULL;
	m_hDsk= NULL;
}
PhysicalDisk::~PhysicalDisk()
{
	free(m_sid);
	CloseHandle(m_hDsk);
}

bool PhysicalDisk::Open(int n)
{
	m_hDsk= OpenDisk(n);
	m_ioctlStyle= GetDiskInfo(m_hDsk, m_info);

	return (m_hDsk!=NULL);
}
bool PhysicalDisk::ReadBlocks(int nStart, int nBlocks, BYTE* &buffer, int& nBytesRead)
{
	return DiskReadBlocks(m_hDsk, nStart, nBlocks, m_info, buffer, nBytesRead);
}
bool PhysicalDisk::WriteBlocks(int nStart, int nBlocks, const BYTE* buffer, int& nBytesWritten)
{
	return DiskWriteBlocks(m_hDsk, nStart, nBlocks, m_info, buffer, nBytesWritten);
}

int PhysicalDisk::blocksize()
{
	return m_info.di_bytes_per_sect;
}

HANDLE OpenDisk(int n)
{
	TCHAR diskname[16];
	_sntprintf(diskname, 16, L"DSK%u:", n);

	HANDLE hDsk= CreateFile( diskname,
			GENERIC_READ|GENERIC_WRITE,
			FILE_SHARE_READ|FILE_SHARE_WRITE,
			NULL,
			OPEN_EXISTING,
			0,
			0 );

	if (hDsk==NULL || hDsk==INVALID_HANDLE_VALUE)
	{
		error("CreateFile('%ls')", diskname);
		return NULL;
	}
	return hDsk;
}


bool DiskReadBlocks(HANDLE hDsk, int nStart, int nBlocks, const DISK_INFO& info,
				BYTE* &buffer, int& nBytesRead)
{
	SG_REQ req;

	if (nBytesRead==0 && buffer==NULL) {
		nBytesRead= nBlocks*info.di_bytes_per_sect;
		buffer= (BYTE*)malloc(nBytesRead);
	}

    req.sr_start= nStart;
	req.sr_num_sec= nBlocks;
	req.sr_num_sg= 1;
    req.sr_status= 0;
    req.sr_callback= NULL;
	req.sr_sglist[0].sb_len= nBytesRead;
	req.sr_sglist[0].sb_buf= buffer;

	DWORD nReturned;
	//if (DeviceIoControl(hDsk, IOCTL_DISK_READ, &req, sizeof(req), NULL, NULL, &nReturned, NULL))
	//	return true;
	if (DeviceIoControl(hDsk, DISK_IOCTL_READ, &req, sizeof(req), NULL, NULL, &nReturned, NULL))
		return true;

	error("DeviceIoControl(IOCTL_DISK_READ)");
	return false;
}
bool DiskWriteBlocks(HANDLE hDsk, int nStart, int nBlocks, const DISK_INFO& info,
			const BYTE* buffer, int& nBytesWritten)
{
	SG_REQ req;

	nBytesWritten= nBlocks*info.di_bytes_per_sect;
    req.sr_start= nStart;
	req.sr_num_sec= nBlocks;
	req.sr_num_sg= 1;
    req.sr_status= 0;
    req.sr_callback= NULL;
	req.sr_sglist[0].sb_len= nBytesWritten;
	req.sr_sglist[0].sb_buf= (BYTE*)buffer;

	DWORD nReturned;
	//if (DeviceIoControl(hDsk, IOCTL_DISK_WRITE, &req, sizeof(req), NULL, NULL, &nReturned, NULL))
	//	return true;
	if (DeviceIoControl(hDsk, DISK_IOCTL_WRITE, &req, sizeof(req), NULL, NULL, &nReturned, NULL))
		return true;

	error("DeviceIoControl(IOCTL_DISK_WRITE)");
	return false;
}

int GetDiskInfo(HANDLE hDsk, DISK_INFO &info)
{
	DWORD nReturned;

	// various ce versions use all possible combinations, so we try 'm all.
	if (DeviceIoControl(hDsk, IOCTL_DISK_GETINFO, &info, sizeof(info), NULL, 0, &nReturned, NULL))
		return 1;
	if (DeviceIoControl(hDsk, IOCTL_DISK_GETINFO, NULL, 0, &info, sizeof(info), &nReturned, NULL))
		return 2;
	if (DeviceIoControl(hDsk, DISK_IOCTL_GETINFO, &info, sizeof(info), NULL, 0, &nReturned, NULL))
		return 3;
	if (DeviceIoControl(hDsk, DISK_IOCTL_GETINFO, NULL, 0, &info, sizeof(info), &nReturned, NULL))
		return 4;

	error("DeviceIoControl(IOCTL_DISK_GETINFO, old/in)");
	return 0;
}
int GetDiskName(HANDLE hDsk, CString &name)
{
	DWORD nReturned;
	if (DeviceIoControl(hDsk, IOCTL_DISK_GETNAME, name.GetBuffer(MAX_PATH), MAX_PATH, NULL, 0, &nReturned, NULL))
		return 1;
	if (DeviceIoControl(hDsk, IOCTL_DISK_GETNAME, NULL, 0, name.GetBuffer(MAX_PATH), MAX_PATH, &nReturned, NULL))
		return 2;
	if (DeviceIoControl(hDsk, DISK_IOCTL_GETNAME, name.GetBuffer(MAX_PATH), MAX_PATH, NULL, 0, &nReturned, NULL))
		return 3;
	if (DeviceIoControl(hDsk, DISK_IOCTL_GETNAME, NULL, 0, name.GetBuffer(MAX_PATH), MAX_PATH, &nReturned, NULL))
		return 4;
	error("DeviceIoControl(IOCTL_DISK_GETNAME, old/out)");
	return 0;
}

int IoctlDiskStorageId(HANDLE hDsk, DWORD nr, bool bAsInParam, STORAGE_IDENTIFICATION* &sid);
int GetDiskStorageId(HANDLE hDsk, STORAGE_IDENTIFICATION* &sid)
{
	if (IoctlDiskStorageId(hDsk, IOCTL_DISK_GET_STORAGEID, true, sid))
		return 1;
	if (IoctlDiskStorageId(hDsk, IOCTL_DISK_GET_STORAGEID, false, sid))
		return 2;
/*
	if (IoctlDiskStorageId(hDsk, DISK_IOCTL_GET_STORAGEID, true, sid))
		return 3;
	if (IoctlDiskStorageId(hDsk, DISK_IOCTL_GET_STORAGEID, false, sid))
		return 4;
*/
	error("DeviceIoControl(IOCTL_DISK_GET_STORAGEID, new/out)");
	return 0;
}
int IoctlDiskStorageId(HANDLE hDsk, DWORD nr, bool bAsInParam, STORAGE_IDENTIFICATION* &sid)
{
	sid= (STORAGE_IDENTIFICATION*)malloc(sizeof(STORAGE_IDENTIFICATION)+64);
	sid->dwSize= sizeof(STORAGE_IDENTIFICATION);
	DWORD nReturned;
	if (DeviceIoControl(hDsk, nr, 
		bAsInParam?sid:NULL, bAsInParam?sizeof(STORAGE_IDENTIFICATION):0,
		bAsInParam?NULL:sid, bAsInParam?0:sizeof(STORAGE_IDENTIFICATION), 
		&nReturned, NULL))
		return true;
	free(sid);
	if (ERROR_INSUFFICIENT_BUFFER != GetLastError())
	{
		sid=NULL;
		return false;
	}

	sid= (STORAGE_IDENTIFICATION*)malloc(sid->dwSize);

	if (DeviceIoControl(hDsk, nr, 
		bAsInParam?sid:NULL, bAsInParam?sizeof(STORAGE_IDENTIFICATION):0,
		bAsInParam?NULL:sid, bAsInParam?0:sizeof(STORAGE_IDENTIFICATION), 
		&nReturned, NULL))
		return true;
	free(sid);

	return false;
}
/*
bool GetDiskDeviceInfo(HANDLE hDsk, STOREAGEDEVICEINFO &di)
{
	DWORD nReturned;
	if (DeviceIoControl(hDsk, IOCTL_DISK_DEVICE_INFO, NULL, 0, &di, sizeof(STOREAGEDEVICEINFO), &nReturned, NULL))
		return true;

	error("DeviceIoControl(IOCTL_DISK_DEVICE_INFO)");
	return false;
}
*/
void DumpDiskInfo(DISK_INFO *info);
void DumpStorageId(STORAGE_IDENTIFICATION *sid);

void PhysicalDisk::DumpInfo()
{
	DISK_INFO info;
	if (GetDiskInfo(m_hDsk, info))
		DumpDiskInfo(&info);

	CString name;
	if (GetDiskName(m_hDsk, name))
		debug("name: '%ls'\n", name.GetBuffer(0));
	
	STORAGE_IDENTIFICATION *sid= NULL;
	if (GetDiskStorageId(m_hDsk, sid))
	{
		DumpStorageId(sid);
		free(sid);
	}
}

void DumpDiskInfo(DISK_INFO *info)
{
	debug("diskinfo: %d sectors, %d bytes, %d cylinders, %d heads, %d sectors per track\n",
		info->di_total_sectors, info->di_bytes_per_sect, info->di_cylinders,
		info->di_heads, info->di_sectors);

	if (info->di_flags&DISK_INFO_FLAG_MBR           ) debug(" has MBR");
	if (info->di_flags&DISK_INFO_FLAG_CHS_UNCERTAIN ) debug(" uncertain CHS");
	if (info->di_flags&DISK_INFO_FLAG_UNFORMATTED   ) debug(" unformatted");
	if (info->di_flags&DISK_INFO_FLAG_PAGEABLE      ) debug(" pageable");

	debug("\n");
}

void DumpStorageId(STORAGE_IDENTIFICATION *sid)
{
	debug("sid: size=%d flags=%08lx mf-ofs=%08lx ser-ofs=%08lx\n", 
			sid->dwSize,
			sid->dwFlags, sid->dwManufactureIDOffset, sid->dwSerialNumOffset);
	if (sid->dwFlags&MANUFACTUREID_INVALID) debug("  mf: %ls\n", ((TCHAR*)sid)+sid->dwManufactureIDOffset);
	if (sid->dwFlags&SERIALNUM_INVALID    ) debug("  sn: %ls\n", ((TCHAR*)sid)+sid->dwSerialNumOffset);
}