/********************************************************/

/*  ntapi: Native API core library                      */

/*  Copyright (C) 2013--2016  Z. Gilboa                 */

/*  Released under GPLv2 and GPLv3; see COPYING.NTAPI.  */

/********************************************************/

#include <ntapi/nt_object.h>

#include <ntapi/nt_file.h>

#include <ntapi/nt_mount.h>

#include <ntapi/nt_atomic.h>

#include <ntapi/ntapi.h>

#include "ntapi_impl.h"

typedef enum __dos_drive_handle_type {

	__DOS_DRIVE_DEVICE_HANDLE,

	__DOS_DRIVE_ROOT_HANDLE

} _dos_drive_handle_type;

typedef struct __dos_name_buffer {

	wchar16_t	global_prefix[4];

	wchar16_t	dos_letter;

	wchar16_t	colon;

	wchar16_t	root;

	wchar16_t	null_termination;

} _dos_name_buffer;

static int32_t __stdcall __tt_connect_to_mount_point_manager(void)

{

	int32_t			status;

	void *			hdev;

	void *			hdev_prev;

	nt_oa			oa;

	nt_iosb			iosb;

	nt_unicode_string	dev_name;

	uint16_t		dev_name_buffer[] = {

				'\\','?','?','\\',

				'M','o','u','n','t',

				'P','o','i','n','t',

				'M','a','n','a','g','e','r',0};

	dev_name.strlen = sizeof(wchar16_t) * (4+5+5+7);

	dev_name.maxlen = 0;

	dev_name.buffer = dev_name_buffer;

	oa.len		= sizeof(nt_oa);

	oa.root_dir	= (void *)0;

	oa.obj_name	= &dev_name;

	oa.obj_attr	= NT_OBJ_CASE_INSENSITIVE;

	oa.sec_desc	= (nt_sd *)0;

	oa.sec_qos	= (nt_sqos *)0;

	status = __ntapi->zw_create_file(

		&hdev,

		NT_SEC_SYNCHRONIZE | NT_FILE_READ_ATTRIBUTES,

		&oa,

		&iosb,

		0,

		NT_FILE_ATTRIBUTE_NORMAL,

		NT_FILE_SHARE_READ | NT_FILE_SHARE_WRITE,

		NT_FILE_OPEN,

		NT_FILE_NON_DIRECTORY_FILE | NT_FILE_SYNCHRONOUS_IO_NONALERT,

		(void *)0,

		0);

	if (status != NT_STATUS_SUCCESS)

		return status;

	hdev_prev = (void *)at_locked_cas(

		(intptr_t *)&__ntapi_internals()->hdev_mount_point_mgr,

		0,(intptr_t)hdev);

	if (hdev_prev)

		__ntapi->zw_close(hdev);

	return status;

}

static int32_t __stdcall __tt_get_dos_drive_device_or_root_handle(

	__out	void **			hdrive,

	__in	wchar16_t *		drive_letter,

	__in	_dos_drive_handle_type	handle_type)

{

	#define		__common_mode	(NT_FILE_SYNCHRONOUS_IO_ALERT)

	#define		__common_access	(NT_SEC_SYNCHRONIZE \

					| NT_FILE_READ_ATTRIBUTES)

	int32_t			status;

	nt_oa			oa;

	nt_iosb			iosb;

	uint32_t		open_flags;

	uint32_t		access_flags;

	nt_unicode_string	dos_name;

	_dos_name_buffer	dos_name_buffer = {

					{'\\','?','?','\\'},

					'_',':',0,0};

	if (!hdrive || !drive_letter)

		return NT_STATUS_INVALID_PARAMETER;

	if ((*drive_letter>='A') && (*drive_letter<='Z'))

		dos_name_buffer.dos_letter = *drive_letter;

	else if ((*drive_letter>='a') && (*drive_letter<='z'))

		dos_name_buffer.dos_letter = *drive_letter + 'A' - 'a';

	else

		return NT_STATUS_INVALID_PARAMETER_2;

	dos_name.strlen = ((size_t)(&((_dos_name_buffer *)0)->root));

	dos_name.maxlen = 0;

	dos_name.buffer = &(dos_name_buffer.global_prefix[0]);

	switch (handle_type) {

		case __DOS_DRIVE_DEVICE_HANDLE:

			open_flags	= __common_mode;

			access_flags	= __common_access;

			break;

		case __DOS_DRIVE_ROOT_HANDLE:

			open_flags	= __common_mode   | NT_FILE_DIRECTORY_FILE;

			access_flags	= __common_access | NT_FILE_READ_ACCESS;

			dos_name_buffer.root = '\\';

			dos_name.strlen += sizeof(wchar16_t);

			break;

		default:

			open_flags	= 0;

			access_flags	= 0;

			break;

	}

	oa.len		= sizeof(nt_oa);

	oa.root_dir	= (void *)0;

	oa.obj_name	= &dos_name;

	oa.obj_attr	= NT_OBJ_INHERIT;

	oa.sec_desc	= (nt_sd *)0;

	oa.sec_qos	= (nt_sqos *)0;

	status = __ntapi->zw_open_file(

		hdrive,

		access_flags,

		&oa,

		&iosb,

		NT_FILE_SHARE_READ | NT_FILE_SHARE_WRITE,

		open_flags);

	return status;

}

int32_t __stdcall	__ntapi_tt_get_dos_drive_device_handle(

	__out	void **			hdevice,

	__in	wchar16_t *		drive_letter)

{

	return __tt_get_dos_drive_device_or_root_handle(

		hdevice,

		drive_letter,

		__DOS_DRIVE_DEVICE_HANDLE);

}

int32_t __stdcall	__ntapi_tt_get_dos_drive_root_handle(

	__out	void **			hroot,

	__in	wchar16_t *		drive_letter)

{

	return __tt_get_dos_drive_device_or_root_handle(

		hroot,

		drive_letter,

		__DOS_DRIVE_ROOT_HANDLE);

}

int32_t __stdcall	__ntapi_tt_get_dos_drive_device_name(

	__in	void *			hdevice		__optional,

	__in	wchar16_t *		drive_letter	__optional,

	__out	nt_mount_dev_name *	buffer,

	__in	uint32_t		buffer_size)

{

	int32_t			status;

	nt_iosb			iosb;

	if (!hdevice && (status = __tt_get_dos_drive_device_or_root_handle(

			&hdevice,

			drive_letter,

			__DOS_DRIVE_DEVICE_HANDLE)))

		return status;

	return __ntapi->zw_device_io_control_file(

		hdevice,

		(void *)0,

		(nt_io_apc_routine *)0,

		(void *)0,

		&iosb,

		NT_IOCTL_MOUNTDEV_QUERY_DEVICE_NAME,

		(void *)0,

		0,

		buffer,

		buffer_size);

}

int32_t __stdcall	__ntapi_tt_get_dos_drive_mount_points(

	__in	void *			hdevice		__optional,

	__in	wchar16_t *		drive_letter	__optional,

	__in	nt_mount_dev_name *	dev_name	__optional,

	__out	void *			buffer,

	__in	uint32_t		buffer_size)

{

	int32_t			status;

	nt_iosb			iosb;

	wchar16_t		dev_name_buffer[64];

	nt_mount_point_param *	dev_mount_point;

	nt_mount_points *	dev_mount_points;

	uintptr_t		addr;

	if (!dev_name) {

		dev_name = (nt_mount_dev_name *)&dev_name_buffer;

		if ((status = __ntapi_tt_get_dos_drive_device_name(

				hdevice,

				drive_letter,

				dev_name,

				sizeof(dev_name_buffer))))

			return status;

	}

	if (buffer_size < sizeof(nt_mount_mgr_mount_point) \

				+ sizeof(nt_mount_dev_name) \

				+ sizeof(dev_name->name_length))

		return NT_STATUS_BUFFER_TOO_SMALL;

	dev_mount_point = (nt_mount_point_param *)buffer;

	dev_mount_point->symlink_name_offset = 0;

	dev_mount_point->symlink_name_length = 0;

	dev_mount_point->unique_id_offset    = 0;

	dev_mount_point->unique_id_length    = 0;

	dev_mount_point->device_name_offset  = ((size_t)(&((nt_mount_point_param *)0)->device_name));

	dev_mount_point->device_name_length  = dev_name->name_length;

	dev_mount_point->mount_points_offset = 0;

	__ntapi->tt_memcpy_utf16(

		dev_mount_point->device_name,

		dev_name->name,

		dev_name->name_length);

	addr = (uintptr_t)(dev_mount_point->device_name) + dev_name->name_length;

	addr += sizeof(uintptr_t) - 1;

	addr /= sizeof(uintptr_t);

	addr *= sizeof(uintptr_t);

	dev_mount_points = (nt_mount_points *)addr;

	if (!__ntapi_internals()->hdev_mount_point_mgr)

		status = __tt_connect_to_mount_point_manager();

	if (!__ntapi_internals()->hdev_mount_point_mgr)

		return status;

	status = __ntapi->zw_device_io_control_file(

		__ntapi_internals()->hdev_mount_point_mgr,

		(void *)0,

		(nt_io_apc_routine *)0,

		(void *)0,

		&iosb,

		NT_IOCTL_MOUNTMGR_QUERY_POINTS,

		dev_mount_point,

		(uint32_t)(uintptr_t)&(((nt_mount_point_param *)0)->device_name) + dev_name->name_length,

		dev_mount_points,

		(uint32_t)((uintptr_t)buffer + buffer_size - addr));

	dev_mount_point->mount_points_offset = (uint16_t)((uintptr_t)addr - (uintptr_t)buffer);

	return status;

}

int32_t __stdcall	__ntapi_tt_dev_mount_points_to_statfs(

	__in		nt_mount_points *	mount_points,

	__in_out	nt_statfs *		statfs)

{

	int32_t		status;

	uint32_t	i;

	nt_mount_mgr_mount_point *	mount_point;

	wchar16_t *			symlink;

	mount_point = mount_points->mount_points;

	statfs->nt_drive_letter = 0;

	for (i = 0; i < mount_points->number; i++, mount_point++) {

		symlink  = (wchar16_t *)mount_points;

		symlink += mount_point->symlink_name_offset / sizeof(wchar16_t);

		if (symlink[0] != '\\')

			return NT_STATUS_UNEXPECTED_IO_ERROR;

		if ((symlink[1] == 'D')

				&& (symlink[2] == 'o')

				&& (symlink[3] == 's')

				&& (symlink[4] == 'D')

				&& (symlink[5] == 'e')

				&& (symlink[6] == 'v')

				&& (symlink[7] == 'i')

				&& (symlink[8] == 'c')

				&& (symlink[9] == 'e')

				&& (symlink[10] == 's'))

			statfs->nt_drive_letter = ((nt_dos_devices_name *)(symlink))->letter;

		else if ((symlink[1] == '?')

				&& (symlink[2] == '?')

				&& (symlink[3] == '\\')

				&& (symlink[4] == 'V')

				&& (symlink[5] == 'o')

				&& (symlink[6] == 'l')

				&& (symlink[7] == 'u')

				&& (symlink[8] == 'm')

				&& (symlink[9] == 'e')

				&& (symlink[10] == '{')) {

			if ((status = __ntapi_tt_utf16_string_to_guid(

					(nt_guid_str_utf16 *)&symlink[10],

					&statfs->nt_volume_guid)))

				return status;

		}

	}

	return NT_STATUS_SUCCESS;

}

int32_t __stdcall	__ntapi_tt_get_dos_drive_letter_from_device(

	__in	void *			hdevice		__optional,

	__out	wchar16_t *		drive_letter,

	__in	nt_mount_dev_name *	dev_name	__optional,

	__out	void *			buffer,

	__in	uint32_t		buffer_size)

{

	int32_t			status;

	wchar16_t		dev_name_buffer[128];

	nt_statfs		statfs;

	uint32_t		offset;

	nt_mount_points *	mnt_points;

	(void)buffer_size;

	if (!dev_name) {

		dev_name = (nt_mount_dev_name *)&dev_name_buffer;

		status = __ntapi_tt_get_dos_drive_device_name(

			hdevice,

			(wchar16_t *)0,

			dev_name,

			sizeof(dev_name_buffer));

		if (status != NT_STATUS_SUCCESS)

			return status;

	}

	offset = ((nt_mount_point_param *)buffer)->mount_points_offset;

	mnt_points = (nt_mount_points *)((uintptr_t)buffer + offset);

	status = __ntapi_tt_dev_mount_points_to_statfs(

		mnt_points,

		&statfs);

	if (status != NT_STATUS_SUCCESS)

		return status;

	*drive_letter = statfs.nt_drive_letter;

	return status;

}