vdkcmd.c

Virtual Disk Driver

			return -1;
		}

		if (!isdigit(**args)) {
			PrintMessage(MSG_INVALID_PART, *args);
			return -1;
		}

		part_number = atol(*(args++));
	}
	else if (isalpha(**args)) {
		drive_letter = (char)toupper(**(args++));
	}
	else {
		PrintMessage(MSG_UNKNOWN_OPTION, *args, "ULINK");
		return -1;
	}

	if (*args) {
		PrintMessage(MSG_USAGE_ULINK);
		return -1;
	}

	if (!drive_letter) {
		// disk, part are specified
		ret = VdkGetDriveLetter(disk_number, part_number, &drive_letter);

		if (ret != ERROR_SUCCESS) {
			PrintMessage(MSG_GET_LINK_NG, disk_number, part_number);
			printf("%s\n", VdkStatusStr(ret));
			return -1;
		}

		if (!isalpha(drive_letter)) {
			PrintMessage(MSG_GET_LINK_NG, disk_number, part_number);
			printf("%s\n", VdkStatusStr(ERROR_FILE_NOT_FOUND));
			return -1;
		}
	}

	if (isalpha(drive_letter)) {
		ret = VdkDelDriveLetter(drive_letter);

		if (ret != ERROR_SUCCESS) {
			PrintMessage(MSG_UNLINK_NG, drive_letter);
			printf("%s\n", VdkStatusStr(ret));
			return -1;
		}
	}

	PrintMessage(MSG_UNLINK_OK, drive_letter);

	return 0;
}

//
//	Print image file information
//	Command Line Parameters:
//		disk number or drive letter
//
int Image(char **args)
{
	ULONG	disk_number;
	ULONG	max_disk;
	DWORD	ret;

	//	ensure that the driver is installed

	if (driver_state == VDK_NOT_INSTALLED) {
		PrintMessage(MSG_NOT_INSTALLED);
		return -1;
	}

	//	ensure that the driver is running

	if (driver_state != SERVICE_RUNNING) {
		PrintMessage(MSG_NOT_RUNNING);
		return -1;
	}

	ret = VdkGetDriverInfo(NULL, &max_disk, NULL, NULL, NULL);

	if (ret != ERROR_SUCCESS) {
		PrintMessage(MSG_GET_CONFIG_NG);
		printf("%s\n", VdkStatusStr(ret));
		return -1;
	}

	if (*args) {
		if (isdigit(**args)) {
			disk_number = atol(*args);
		}
		else if (isalpha(**args)) {
			ret = GetDiskFromDrive(**args, &disk_number);

			if (ret != ERROR_SUCCESS) {
				PrintMessage(MSG_RESOLVE_LINK_NG, **args);
				return -1;
			}
		}
		else {
			PrintMessage(MSG_INVALID_DISK, *args);
			return -1;
		}

		if (disk_number >= max_disk) {
			PrintMessage(MSG_INVALID_DISK, *args);
			return -1;
		}

		max_disk = disk_number;
	}
	else {
		disk_number = 0;
	}

	do {
		HANDLE hDisk;
		PVDK_OPEN_FILE_INFO	file_info;

		PrintMessage(MSG_VIRTUAL_DISK, disk_number);

		hDisk = VdkOpenDevice(disk_number, 0);

		if (hDisk == INVALID_HANDLE_VALUE) {
			ret = VdkLastError();
			PrintMessage(MSG_GET_IMAGE_NG, disk_number);
			printf("%s\n", VdkStatusStr(ret));
			continue;
		}

		ret = VdkCheckDeviceState(hDisk, 0, 0);

		if (ret != ERROR_SUCCESS) {
			if (ret == ERROR_BUSY) {
				PrintMessage(MSG_DEVICE_BUSY);
				printf("\n");
			}
			else if (ret == ERROR_NOT_READY) {
				PrintMessage(MSG_IMAGE_NONE);
				printf("\n");
			}
			else {
				printf("%s\n\n", VdkStatusStr(ret));
			}

			CloseHandle(hDisk);
			continue;
		}

		ret = VdkGetFileInfo(hDisk, 0, &file_info);

		if (ret != ERROR_SUCCESS) {
			CloseHandle(hDisk);
			PrintMessage(MSG_GET_IMAGE_NG, disk_number);
			printf("%s\n", VdkStatusStr(ret));
			continue;
		}

		//	print image file information

		PrintFileInfo(file_info);

		if (!file_info) {
			CloseHandle(hDisk);
			continue;
		}

		if (!file_info->DiskType) {
			VdkFreeMem(file_info);
			CloseHandle(hDisk);
			continue;
		}

		//
		//	Print partition information
		//
		PrintMessage(MSG_PARTITION_HEADER);

		ret = VdkListPartitions(NULL, hDisk, file_info->Capacity, PartList_Callback, (PVOID)disk_number);

		VdkFreeMem(file_info);
		CloseHandle(hDisk);

		if (ret != ERROR_SUCCESS) {
			PrintMessage(MSG_GET_PART_NG);
			printf("%s\n", VdkStatusStr(ret));
			continue;
		}

		printf("\n");
	}
	while (++disk_number < max_disk);

	return 0;
}

//
//	Print usage help
//	Command Line Parameters:
//		(optional) command
//
int	Help(char **args)
{
	DWORD msg = MSG_HELP_GENERAL;

	if (args && *args) {
		int idx = 0;

		while (Commands[idx].cmd) {
			if (_stricmp(*args, Commands[idx].cmd) == 0) {
				msg = Commands[idx].helpmsg;
				break;
			}
			idx++;
		}

		if (!Commands[idx].cmd) {
			PrintMessage(MSG_UNKNOWN_COMMAND, *args);
			msg = MSG_HELP_HELP;
		}
	}

	PrintMessage(msg);

	return 0;
}

//
//	Print device information
//
int	Device(char **args)
{
	ULONG device_num;
	PVDK_DEVICE_INFO device_info;

	UNREFERENCED_PARAMETER(args);

	if (VdkGetDeviceList(&device_num, &device_info) != ERROR_SUCCESS) {
		return -1;
	}

	while (device_num--) {
		printf("Device Type  : %s\n", (device_info[device_num].DeviceType == VDK_DEVICE_DISK) ? "DISK" : "PART");
		printf("Device Name  : %s\n", device_info[device_num].DeviceName);
		printf("Symbolic Link: %s\n", device_info[device_num].SymbolicLink);
		printf("Reference    : %lu\n", device_info[device_num].ReferenceCount);
		printf("\n");
	}

	VdkFreeMem(device_info);
	return 0;
}

#ifdef _DEBUG

struct _traceflags {
	char *name;
	ULONG value;
}
flagtable[] = {
	"WARN",		VDKWARN,
	"INFO",		VDKINFO,
	"OPEN",		VDKOPEN,
	"CLOSE",	VDKCLOSE,
	"READ",		VDKREAD,
	"WRITE",	VDKWRITE,
	"DISPATCH",	VDKDISPATCH,
	"IOCTL",	VDKIOCTL,
	"UPDATE",	VDKUPDATE,
	"CREATE",	VDKCREATE,
	"DELETE",	VDKDELETE,
	"FORMAT",	VDKFORMAT,
	"SERVER",	VDKSERVER,
	"CLIENT",	VDKCLIENT,
	NULL, 0
};

ULONG ParseFlag(PCHAR name)
{
	int idx;

	if (!_stricmp(name, "ALL")) {
		return 0xffffffff;
	}

	idx = 0;

	while (flagtable[idx].name) {
		if (!_stricmp(name, flagtable[idx].name)) {
			return flagtable[idx].value;
		}
		idx++;
	}

	return 0;
}

int	Trace(char **args)
{
	ULONG	flags, changed;
	int idx;

	flags = VdkTraceFlags(NULL);
	changed = flags;

	while (*args) {
		if (**args == '-') {
			changed &= ~(ParseFlag(*args + 1));
		}
		else if (**args == '+') {
			changed |= ParseFlag(*args + 1);
		}
		else {
			changed |= ParseFlag(*args);
		}
		args++;
	}

	if (changed != flags) {
		flags = VdkTraceFlags(&changed);
	}

	printf("Trace flag is : 0x%08x\n", flags);

	printf("Enabled : ");

	idx = 0;

	while (flagtable[idx].name) {
		if ((flagtable[idx].value & flags) == flagtable[idx].value) {
			printf("%s ", flagtable[idx].name);
		}
		idx++;
	}

	printf("\nDisabled: ");

	idx = 0;

	while (flagtable[idx].name) {
		if ((flagtable[idx].value & flags) != flagtable[idx].value) {
			printf("%s ", flagtable[idx].name);
		}
		idx++;
	}
	printf("\n");
	return 0;
}
#endif _DEBUG

//
//	Get disk number from a drive letter
//
DWORD GetDiskFromDrive(CHAR drive, PULONG disk)
{
	TCHAR	dos_device[] = " :";
	TCHAR	device_name[MAX_DEVNAME_LEN], dos_name[MAX_PATH];
	DWORD	ret;

	dos_device[0] = drive;

	if (!QueryDosDevice(dos_device, dos_name, sizeof(dos_name))) {
		return GetLastError();
	}

	ret = VdkGetDriverInfo(NULL, disk, NULL, NULL, NULL);

	if (ret != ERROR_SUCCESS) {
		return ret;
	}

	while ((*disk)--) {
		ret = VdkGetDeviceName((*disk), (ULONG)-1, device_name);

		if (ret != ERROR_SUCCESS) {
			return ret;
		}

		if (!_strnicmp(dos_name, device_name, strlen(device_name))) {
			return ERROR_SUCCESS;
		}
	}

	return ERROR_FILE_NOT_FOUND;
}

//
//	Print image file information
//
void PrintFileInfo(PVDK_OPEN_FILE_INFO file_info)
{
	ULONG total_sectors;
	PCHAR name_pos;
	ULONG idx;

	if (!file_info) {
		PrintMessage(MSG_IMAGE_NONE);
		printf("\n");
		return;
	}

	//	Virtual disk access type

	if (file_info->DiskType) {
		PrintMessage(MSG_ACCESS_TYPE);

		switch (file_info->DiskType) {
		case VDK_DISKTYPE_WRITABLE:
			PrintMessage(MSG_ACCESS_RW);
			break;

		case VDK_DISKTYPE_READONLY:
			PrintMessage(MSG_ACCESS_RO);
			break;

		case VDK_DISKTYPE_WRITEBLOCK:
			PrintMessage(MSG_ACCESS_WB);
			break;
		}
	}

	PrintMessage(MSG_DISK_CAPACITY,
		file_info->Capacity,
		file_info->Capacity / 2048);

	if (file_info->Cylinders &&
		file_info->Tracks &&
		file_info->Sectors) {

		PrintMessage(MSG_DISK_GEOMETRY,
			file_info->Cylinders,
			file_info->Tracks,
			file_info->Sectors);
	}

	PrintMessage(MSG_DISK_FILES,
		file_info->FilesTotal);

	PrintMessage(MSG_FILE_HEADER);

	total_sectors = 0;
	name_pos = (PCHAR)&(file_info->Files[file_info->FilesTotal]);

	for (idx = 0; idx < file_info->FilesTotal; idx++) {

		if (total_sectors == 0) {
			printf(" -------  -------  ----\n");
		}
		switch (file_info->Files[idx].FileType) {
		case VDK_FILETYPE_NONE:
			printf("  NONE ");
			break;

		case VDK_FILETYPE_FLAT:
			printf("  FLAT ");
			break;

		case VDK_FILETYPE_COWD:
			printf("  COWD ");
			break;

		case VDK_FILETYPE_VMDK:
			printf("  VMDK ");
			break;

		default:
			printf("  ???  ");
			break;
		}

		printf("  %8lu  %s\n",
			file_info->Files[idx].Capacity,
			name_pos);

		name_pos += file_info->Files[idx].NameLength;

		total_sectors += file_info->Files[idx].Capacity;

		if (total_sectors == file_info->Capacity) {
			total_sectors = 0;
		}
	}
	printf("\n");
}

//
//	Confirms unmount retry
//
DWORD Retry_Callback(PVOID param, DWORD err)
{
	printf("\n");
	PrintMessage(MSG_DISMOUNT_NG);
	printf("%s\n", VdkStatusStr(err));

	if (param) {
		// quiet / force mode
		return FALSE;
	}
	else {
		// normal mode
		PrintMessage(MSG_DRIVE_IN_USE);
		return (InputChar(MSG_PROMPT_RETRY, "yn") == 'y');
	}
}

//
//	Confirms unmount continue
//
DWORD Continue_Callback(PVOID param, DWORD err)
{
	UNREFERENCED_PARAMETER(err);

	if ((ULONG)param == CLOSE_QUIET) {
		// quiet mode
		return FALSE;
	}
	else if ((ULONG)param == CLOSE_FORCE) {
		// force mode
		PrintMessage(MSG_CLOSE_FORCED);
		return TRUE;
	}
	else {
		// prompt mode
		printf("\n");
		PrintMessage(MSG_FORCING_WARN);

		return (InputChar(MSG_PROMPT_CONTINUE, "yn") == 'y');
	}
}

//
//	VdkListPartitions callback
//	called for each partition item
//
void PartList_Callback(PPARTITION_ITEM pitem, PVOID param)
{
	CHAR drive[] = " :";

	if ((ULONG)param != (ULONG)-1) {
		VdkGetDriveLetter(
			(ULONG)param,
			pitem->idx,
			&drive[0]);
	}

	if (!isalpha(drive[0])) {
		strcpy(drive, "  ");
	}

	if (pitem->idx) {
		PCHAR name = GetPartitionTypeName(pitem->type);

		printf(" %s  %2d%c  %12lu   %8lu (%6lu MB)  %02xh:%s\n",
			drive,
			pitem->idx,
			pitem->num > 4 ? '*' : ' ',
			pitem->offset,
			pitem->length,
			pitem->length / 2048,
			pitem->type,
			name ? name : "");
	}
	else {
		printf(" %s  %2d%c  %12lu   %8lu (%6lu MB)  %s\n",
			drive,
			pitem->idx,
			pitem->num > 4 ? '*' : ' ',
			pitem->offset,
			pitem->length,
			pitem->length / 2048,
			pitem->fsname);
	}
}

//
//	VdkGetDriveLayout callback
//	called for each recognized partition item
//
void Assign_Callback(PPARTITION_ITEM pitem, PVOID param)
{
	PASSIGN_PARAM assign = (PASSIGN_PARAM)param;
	VDKSTAT ret;

	if ((assign->part_number == (ULONG)-1 &&
		IsPartitionMountable(pitem->type, assign->read_only)) ||
		assign->part_number == pitem->idx) {

		CHAR letter = ChooseDriveLetter();

		if (!isalpha(letter)) {
			PrintMessage(MSG_DRIVE_FULL);
			goto print_info;
		}

		if (assign->drive_letters &&
			isalpha(*(assign->drive_letters)) &&
			*(assign->drive_letters) > letter) {

			letter = *((assign->drive_letters)++);
		}

		ret = VdkSetDriveLetter(
			assign->disk_number,
			pitem->idx,
			letter);

		if (ret != ERROR_SUCCESS) {
			PrintMessage(MSG_LINK_NG, letter, assign->disk_number, pitem->idx);
			printf("%s\n", VdkStatusStr(ret));
		}
	}

print_info:
	//
	//	print partition information
	//
	PartList_Callback(pitem, (PVOID)(assign->disk_number));
}

//
//	Check if current user belongs to Administrators group
//
BOOL IsUserAdmin()
{
	DWORD	i, dwSize = 0;
	HANDLE	hToken;
	PTOKEN_GROUPS pGroupInfo;
	BYTE	sidBuffer[100];
	PSID	pSID = (PSID)&sidBuffer;
	SID_IDENTIFIER_AUTHORITY SIDAuth = SECURITY_NT_AUTHORITY;
	BOOL	ret = FALSE;

	// Open a handle to the access token for the calling process.

	if (!OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &hToken)) {
		return FALSE;
	}

	// Call GetTokenInformation to get the buffer size.

	if(!GetTokenInformation(hToken, TokenGroups, NULL, dwSize, &dwSize)) {

		if(GetLastError() != ERROR_INSUFFICIENT_BUFFER ) {
			return FALSE;
		}
	}

	// Allocate the buffer.

	pGroupInfo = (PTOKEN_GROUPS)GlobalAlloc(GPTR, dwSize);

	// Call GetTokenInformation again to get the group information.

	if (!GetTokenInformation(hToken, TokenGroups, pGroupInfo, dwSize, &dwSize)) {
		goto cleanup;
	}

	// Create a SID for the BUILTIN\Administrators group.

	if (!AllocateAndInitializeSid(
		&SIDAuth, 2, SECURITY_BUILTIN_DOMAIN_RID,
		DOMAIN_ALIAS_RID_ADMINS, 0, 0, 0, 0, 0, 0, &pSID)) {
		goto cleanup;
	}

	// Loop through the group SIDs looking for the administrator SID.

	for (i = 0; i < pGroupInfo->GroupCount; i++) {

		if (EqualSid(pSID, pGroupInfo->Groups[i].Sid)) {
			ret = TRUE;
			break;
		}
	}

cleanup:
	if (pSID) {
		FreeSid(pSID);
	}

	if (pGroupInfo) {
		GlobalFree(pGroupInfo);
	}

	return ret;
}