/*=============================================================================
	FILE			 : sector_io.c 
	ABSTRACT		 : User application sending data to kernel driver which reads/writes
					directly on disk sectors
	AUTHOR		 : Deepak Gupta
     
	Thanks to Chew Keong TAN because some of the driver laoding and unloading functions
	are taken frome there
   				
=============================================================================*/
// sector_io.c : Defines the entry point for the console application.
//

#include "sector_io.h"

#define	DRV_NAME	"sectorio"
#define DRV_FILENAME	"sectorio.sys"

#define SECTOR_IO_DEVICE       0x8000

#define IOCTL_SECTOR_READ		CTL_CODE(SECTOR_IO_DEVICE, 0x800, METHOD_BUFFERED, FILE_READ_ACCESS | FILE_WRITE_ACCESS)
#define IOCTL_SECTOR_WRITE		CTL_CODE(SECTOR_IO_DEVICE, 0x801, METHOD_BUFFERED, FILE_READ_ACCESS | FILE_WRITE_ACCESS)
#define IOCTL_GET_SECTOR_SIZE	CTL_CODE(SECTOR_IO_DEVICE, 0x802, METHOD_BUFFERED, FILE_READ_ACCESS | FILE_WRITE_ACCESS)

BOOL getLoadDriverPriv()
{
	HANDLE hToken;
	LUID huid;
	LUID_AND_ATTRIBUTES priv;
	TOKEN_PRIVILEGES tp;

	if(OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY | TOKEN_ADJUST_PRIVILEGES, &hToken))
	{

		if(LookupPrivilegeValueA(NULL, "SeLoadDriverPrivilege", &huid))
		{
			priv.Attributes = SE_PRIVILEGE_ENABLED;
			priv.Luid = huid;
			
			tp.PrivilegeCount = 1;
			tp.Privileges[0] = priv;

			if(AdjustTokenPrivileges(hToken, FALSE, &tp, 0, NULL, NULL))
			{
				return TRUE;
			}
		}
	}
	return FALSE;
}

BOOL setupRegistry()
{
	HKEY hkey;
	DWORD val;
	char *imgName = "System32\\DRIVERS\\"DRV_FILENAME;

	if(RegCreateKeyA(HKEY_LOCAL_MACHINE, "System\\CurrentControlSet\\Services\\"DRV_NAME, &hkey) != ERROR_SUCCESS)
		return FALSE;

	
	val = 1;
	if(RegSetValueExA(hkey, "Type", 0, REG_DWORD, (PBYTE)&val, sizeof(val)) != ERROR_SUCCESS)
		return FALSE;

	if(RegSetValueExA(hkey, "ErrorControl", 0, REG_DWORD, (PBYTE)&val, sizeof(val)) != ERROR_SUCCESS)
		return FALSE;
	
	val = 3;
	if(RegSetValueExA(hkey, "Start", 0, REG_DWORD, (PBYTE)&val, sizeof(val)) != ERROR_SUCCESS)
		return FALSE;

	
	if(RegSetValueExA(hkey, "ImagePath", 0, REG_EXPAND_SZ, (PBYTE)imgName, strlen(imgName)) != ERROR_SUCCESS)
		return FALSE;

	return TRUE;
}

BOOL loadDriver()
{
	// call ntdll APIs
	HMODULE hntdll;
	ANSI_STRING aStr;
	UNICODE_STRING uStr;

	NTSTATUS (WINAPI * _RtlAnsiStringToUnicodeString)
		(PUNICODE_STRING  DestinationString,
		 IN PANSI_STRING  SourceString,
		 IN BOOLEAN);

	VOID (WINAPI *_RtlInitAnsiString)
		(IN OUT PANSI_STRING  DestinationString,
		 IN PCHAR  SourceString);

	NTSTATUS (WINAPI * _ZwLoadDriver)
		(IN PUNICODE_STRING DriverServiceName);

	NTSTATUS (WINAPI * _ZwUnloadDriver)
		(IN PUNICODE_STRING DriverServiceName);

	VOID (WINAPI * _RtlFreeUnicodeString)
		(IN PUNICODE_STRING  UnicodeString);


	hntdll = GetModuleHandleA("ntdll.dll");
			
	*(FARPROC *)&_ZwLoadDriver = GetProcAddress(hntdll, "NtLoadDriver");

	*(FARPROC *)&_ZwUnloadDriver = GetProcAddress(hntdll, "NtUnloadDriver");
				
	*(FARPROC *)&_RtlAnsiStringToUnicodeString = 
			GetProcAddress(hntdll, "RtlAnsiStringToUnicodeString");

	*(FARPROC *)&_RtlInitAnsiString = 
			GetProcAddress(hntdll, "RtlInitAnsiString");

	*(FARPROC *)&_RtlFreeUnicodeString = 
			GetProcAddress(hntdll, "RtlFreeUnicodeString");

	if(_ZwLoadDriver && _ZwUnloadDriver && _RtlAnsiStringToUnicodeString &&
	   _RtlInitAnsiString && _RtlFreeUnicodeString)
	{
		_RtlInitAnsiString(&aStr, 
		"\\Registry\\Machine\\System\\CurrentControlSet\\Services\\"DRV_NAME);
						
		if(_RtlAnsiStringToUnicodeString(&uStr, &aStr, TRUE) != STATUS_SUCCESS)
			return FALSE;
		else
		{
			if(_ZwLoadDriver(&uStr) == STATUS_SUCCESS)
			{
				_RtlFreeUnicodeString(&uStr);
				return TRUE;
			}
			_RtlFreeUnicodeString(&uStr);
		}
	}

	return FALSE;
}

void cleanupDriver(void)
{
	char sysDir[MAX_PATH + 1];
	GetSystemDirectoryA(sysDir, MAX_PATH);
	strncat(sysDir, "\\drivers\\"DRV_FILENAME, MAX_PATH);
	DeleteFileA(sysDir);

	RegDeleteKeyA(HKEY_LOCAL_MACHINE, "System\\CurrentControlSet\\Services\\"DRV_NAME"\\Enum");
	RegDeleteKeyA(HKEY_LOCAL_MACHINE, "System\\CurrentControlSet\\Services\\"DRV_NAME);
}

BOOL unloadDriver()
{
	// call ntdll APIs
	HMODULE hntdll;
	ANSI_STRING aStr;
	UNICODE_STRING uStr;

	NTSTATUS (WINAPI * _RtlAnsiStringToUnicodeString)
		(PUNICODE_STRING  DestinationString,
		 IN PANSI_STRING  SourceString,
		 IN BOOLEAN);

	VOID (WINAPI *_RtlInitAnsiString)
		(IN OUT PANSI_STRING  DestinationString,
		 IN PCHAR  SourceString);

	NTSTATUS (WINAPI * _ZwLoadDriver)
		(IN PUNICODE_STRING DriverServiceName);

	NTSTATUS (WINAPI * _ZwUnloadDriver)
		(IN PUNICODE_STRING DriverServiceName);

	VOID (WINAPI * _RtlFreeUnicodeString)
		(IN PUNICODE_STRING  UnicodeString);


	hntdll = GetModuleHandleA("ntdll.dll");
			
	*(FARPROC *)&_ZwLoadDriver = GetProcAddress(hntdll, "NtLoadDriver");

	*(FARPROC *)&_ZwUnloadDriver = GetProcAddress(hntdll, "NtUnloadDriver");
				
	*(FARPROC *)&_RtlAnsiStringToUnicodeString = 
			GetProcAddress(hntdll, "RtlAnsiStringToUnicodeString");

	*(FARPROC *)&_RtlInitAnsiString = 
			GetProcAddress(hntdll, "RtlInitAnsiString");

	*(FARPROC *)&_RtlFreeUnicodeString = 
			GetProcAddress(hntdll, "RtlFreeUnicodeString");

	if(_ZwLoadDriver && _ZwUnloadDriver && _RtlAnsiStringToUnicodeString &&
	   _RtlInitAnsiString && _RtlFreeUnicodeString)
	{

		_RtlInitAnsiString(&aStr, 
		"\\Registry\\Machine\\System\\CurrentControlSet\\Services\\"DRV_NAME);
						
		if(_RtlAnsiStringToUnicodeString(&uStr, &aStr, TRUE) != STATUS_SUCCESS)
			return FALSE;
		else
		{
			if(_ZwUnloadDriver(&uStr) == STATUS_SUCCESS)
			{
				_RtlFreeUnicodeString(&uStr);
				return TRUE;
			}
			_RtlFreeUnicodeString(&uStr);
		}
	}

	return FALSE;
}


void uninstallDriver(void)
{
	char drvFullPath[MAX_PATH+1];
	char *filePart;
	HANDLE hFile;
	char sysDir[MAX_PATH + 1];

	ZeroMemory(drvFullPath, MAX_PATH);		
	GetFullPathNameA(DRV_FILENAME, MAX_PATH, drvFullPath, &filePart);

	hFile= CreateFileA(drvFullPath, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING,
								FILE_ATTRIBUTE_NORMAL, 0);
	if(hFile == INVALID_HANDLE_VALUE)
	{
		printf("Cannot find required driver file %s\n", drvFullPath);
		return;
	}
	else
	{
		CloseHandle(hFile);

		GetSystemDirectoryA(sysDir, MAX_PATH);
		strncat(sysDir, "\\drivers\\"DRV_FILENAME, MAX_PATH);
		CopyFileA(drvFullPath, sysDir, TRUE);

		if(!getLoadDriverPriv())
		{
			printf("Error getting load driver privilege!\n");
		}
		else
		{
			if(!setupRegistry())
			{
				printf("Error setting driver registry keys!\nMake sure you are running this as Administrator.\n");
			}
			else
			{
				if(unloadDriver())
					printf("Support driver successfully unloaded.\n");
				else
					printf("Unload support driver failed.  It is probably not loaded.\n");
			}
		}
		cleanupDriver();
	}
}

HANDLE openDriver(void)
{
	HANDLE hDevice;
	HANDLE hFile;
	char drvFullPath[MAX_PATH+1];
	char *filePart;
	char sysDir[MAX_PATH + 1];

	hDevice = CreateFileA("\\\\.\\"DRV_NAME, GENERIC_WRITE | GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE,
		   				 NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);

	if(hDevice == INVALID_HANDLE_VALUE)
	{		


		ZeroMemory(drvFullPath, MAX_PATH);		
		GetFullPathNameA(DRV_FILENAME, MAX_PATH, drvFullPath, &filePart);
		
		hFile= CreateFileA(drvFullPath, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING,
									FILE_ATTRIBUTE_NORMAL, 0);
		if(hFile == INVALID_HANDLE_VALUE)
		{
			printf("Cannot find required driver file %s\n", drvFullPath);
			return INVALID_HANDLE_VALUE;
		}
		else
		{
			CloseHandle(hFile);

			GetSystemDirectoryA(sysDir, MAX_PATH);
			strncat(sysDir, "\\drivers\\"DRV_FILENAME, MAX_PATH);
			CopyFileA(drvFullPath, sysDir, TRUE);

			if(!getLoadDriverPriv())
			{
				printf("Error getting load driver privilege!\n");
			}
			else
			{
				if(!setupRegistry())
				{
					printf("Error setting driver registry keys!\nMake sure you are running this as Administrator.\n");
				}
				else
				{
					loadDriver();
					hDevice = CreateFileA("\\\\.\\"DRV_NAME, GENERIC_WRITE | GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE,
		   				 NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
					if(hDevice == INVALID_HANDLE_VALUE)
					{
						printf("Error loading kernel support driver!\nMake sure you are running this as Administrator.\n");
					}
				}
			}
			cleanupDriver();
		}
	}
	
	return hDevice;
}

DWORD DoDeviceIoCtl(DWORD dwIoCtl, DWORD dwDiskObjOrdinal, 
					BOOLEAN bIsRawDisk, ULONGLONG ullSectorNumber, PVOID *pBuf)
{
	CHAR	szDriverPath[MAX_PATH];
	HANDLE	hDevice;
	DWORD	dwStatus, Size = 512, Bytes;
	DISK_LOCATION	dlInfo;
	PVOID	pMem = NULL;
	PBYTE	pByte;
	OFSTRUCT	ofReOpenBuff;
	DWORD		i;
	char	c;

	hDevice = openDriver();

	if (hDevice == INVALID_HANDLE_VALUE) {
		printf("Can't Open Driver\n");
	}

	dlInfo.bIsRawDiskObj = bIsRawDisk;
	dlInfo.dwDiskOrdinal = dwDiskObjOrdinal;
	dlInfo.ullSectorNum  = ullSectorNumber;
	

	do
	{
		if (dwIoCtl == IOCTL_SECTOR_WRITE) {
			Size += sizeof(DISK_LOCATION);
		}

		pMem = malloc(Size);

		// A very ugly hack to transfer disk location data and input buffer both for write operations
		// Came to know about bug of write operations very late, So instead of handling mapping user address into kernel
		// address space, I did this hack, Will fix it in future
	
		if (!pMem)
			return ERROR_NOT_ENOUGH_MEMORY;

		memset(pMem, 0x00, Size);
		
		if (dwIoCtl == IOCTL_SECTOR_WRITE) {
			pByte = (PBYTE) ((ULONG)pMem + sizeof(DISK_LOCATION));
			i = 0;
			printf("Please type the data you want to write to the sector (input redirection is a better option)\nInput Data: \n");

			do{
				c = getc(stdin);
				pByte[i] = c;
				i++;
			}while ((i < Size) && (c != EOF));
		}
		if (dwIoCtl == IOCTL_SECTOR_WRITE) {
			memcpy(pMem, &dlInfo, sizeof(dlInfo));
			if(DeviceIoControl(hDevice, dwIoCtl, pMem, Size, NULL, 0, &Bytes, NULL))
				break;
		} else {
			if(DeviceIoControl(hDevice, dwIoCtl, &dlInfo, sizeof(dlInfo), pMem, Size, &Bytes, NULL))
				break;
		}
		
		dwStatus = GetLastError();
		printf ("DeviceIoControl Failed and error code is %d\n", dwStatus);

		if (dwIoCtl == IOCTL_SECTOR_WRITE) {
			Size -= sizeof(DISK_LOCATION);
		}
		free(pMem);
		pMem = NULL;
		Size = Size * 2;
	}while(dwStatus == ERROR_INSUFFICIENT_BUFFER);

	if (dwIoCtl == IOCTL_SECTOR_READ && pMem) {
		pByte = (PBYTE) pMem;
		i = 0;
		c = 0;
		printf("Displaying the data read from the sector (in hexadecimal, output redirection can also work)\nOutput Data: \n");
		do{
			c = pByte[i];
			//putc(c, stdout);
			printf("0x%-02X ", c & 0x00FFUL);

			if (!(i+1)%0x10) {
				printf ("\n");
			}
			i++;
		}while (i < Size);
	}

	*pBuf = pMem;
	return Size;
}
VOID PrintUsage()
{
	printf("Usage is:\n"
		"DiskSector {/disk | /partition} <rawdisk number | partition number> " 
		"{/read | /write} <sectornumber> {/unload}\n"
		"\n{/disk | /partition} <rawdisk number | partition number>\n"
		"Disk and Parition options are mutually exclusive\n"
		"Disk numbering starts from 0 while partition starts from 1\n"
		"\n{/read | /write} <sectornumber>\n"
		"Read and Write options are mutually exclusive\n"
		"Sector numbering starts from 0\n"
		"\n{/unload} \nThis option simply unloads the support driver\n"
		"\ne.g \"DiskSector /disk 0 /read 0\" will read raw sector 0 of harddisk 0\n");
	return;
}

int __cdecl main(int argc, char* argv[])
{
	int			argIndex;
	BOOLEAN		bIsRawDisk, bReadWrite, bLoadDriver = TRUE;
	DWORD		dwDiskObjOrdinal = -1;
	ULONGLONG	ullSectorNumber = -1;
	DWORD		dwSize = 512;
	PVOID		pBuf = NULL;

	
	if( argc < 2 || strcmp(argv[1], "/?") == 0 ) {
        PrintUsage();
        return (0);
    }

	argIndex = 1;

	for (argIndex = 1; argIndex < argc; argIndex++)
	{
		if (!strcmp(argv[argIndex], "/disk") || !strcmp(argv[argIndex], "/partition")) {
			if (dwDiskObjOrdinal == -1) {
				bIsRawDisk = strcmp(argv[argIndex], "/disk") ? FALSE : TRUE;
				argIndex++;
				if (argIndex < argc) {
					char *endptr;
					dwDiskObjOrdinal = strtoul(argv[argIndex], &endptr, 10);
				} else {
					PrintUsage();
					return -1;
				}
			} else {
				PrintUsage();
				return -1;
			}
		} else if (!strcmp(argv[argIndex], "/read") || !strcmp(argv[argIndex], "/write")) {
			if (ullSectorNumber == -1) {
				bReadWrite = strcmp(argv[argIndex], "/read") ? FALSE : TRUE;
				argIndex++;
				if (argIndex < argc) {
					char *endptr;
					ullSectorNumber = _strtoui64(argv[argIndex], &endptr, 10);
				}
			} else {
				PrintUsage();
				return -1;
			}

		} else if (!strcmp(argv[argIndex], "/unload")) {
			bLoadDriver = FALSE;
		} else {
			PrintUsage();
			return -1;
		}
	}

	if (!bLoadDriver) {
		uninstallDriver();
		return 0;
	} else if (dwDiskObjOrdinal == -1 || ullSectorNumber == -1){
		PrintUsage();
		return -1;

	}

	DoDeviceIoCtl(bReadWrite?IOCTL_SECTOR_READ:IOCTL_SECTOR_WRITE, 
				dwDiskObjOrdinal, bIsRawDisk, ullSectorNumber, &pBuf);
	if (!pBuf)
		free(pBuf);

	return 0;
}